# -*- coding: binary -*-

require 'rex/proto/http/web_socket'
require 'uri'

module Msf
  class Exploit
    class Remote
      module HTTP
        module Kubernetes
          class Client
            USER_AGENT = 'kubectl/v1.22.2 (linux/amd64) kubernetes/8b5a191'.freeze

            class ExecChannel < Rex::Proto::Http::WebSocket::Interface::Channel
              attr_reader :error

              def initialize(websocket)
                @error = {}

                super(websocket, write_type: :text)
              end

              def on_data_read(data, _data_type)
                return data if data.blank?

                exec_channel = data[0].ord
                data = data[1..-1]
                case exec_channel
                when EXEC_CHANNEL_STDOUT
                  return data
                when EXEC_CHANNEL_STDERR
                  return data
                when EXEC_CHANNEL_ERROR
                  @error = JSON(data)
                end

                nil
              end

              def on_data_write(data)
                EXEC_CHANNEL_STDIN.chr + data
              end
            end

            def initialize(config)
              @http_client = config.fetch(:http_client)
              @token = config[:token]
            end

            # rubocop:disable Style/DoubleNegation
            def exec_pod(name, namespace, command, options = {})
              options = {
                'stdin' => false,
                'stdout' => false,
                'stderr' => false,
                'tty' => false
              }.merge(options)

              # while kubectl uses SPDY/3.1, the Python client uses WebSockets over HTTP/1.1
              # see: https://github.com/kubernetes/kubernetes/issues/7452
              websocket = http_client.connect_ws(
                request_options(
                  {
                    'method' => 'GET',
                    'uri' => http_client.normalize_uri("/api/v1/namespaces/#{namespace}/pods/#{name}/exec"),
                    'query' => URI.encode_www_form(
                      {
                        'command' => command,
                        'stdin' => !!options.delete('stdin'),
                        'stdout' => !!options.delete('stdout'),
                        'stderr' => !!options.delete('stderr'),
                        'tty' => !!options.delete('tty')
                      }
                    ),
                    'headers' => {
                      'Sec-Websocket-Protocol' => 'v4.channel.k8s.io'
                    }
                  },
                  options
                )
              )

              websocket
            end
            # rubocop:enable Style/DoubleNegation

            def exec_pod_capture(name, namespace, command, options = {}, &block)
              websocket = exec_pod(name, namespace, command, options)
              return nil if websocket.nil?

              result = { error: {}, stdout: '', stderr: '' }
              websocket.wsloop do |channel_data, _data_type|
                next if channel_data.blank?

                channel = channel_data[0].ord
                channel_data = channel_data[1..-1]
                case channel
                when EXEC_CHANNEL_STDOUT
                  result[:stdout] << channel_data
                  block.call(channel_data, nil) if block_given?
                when EXEC_CHANNEL_STDERR
                  result[:stderr] << channel_data
                  block.call(nil, channel_data) if block_given?
                when EXEC_CHANNEL_ERROR
                  result[:error] = JSON(channel_data)
                end
              end

              result
            end

            def get_version(options = {})
              _res, json = call_api(
                {
                  'method' => 'GET',
                  'uri' => http_client.normalize_uri("/version")
                },
                options
              )

              json
            end

            def get_pod(pod, namespace, options = {})
              _res, json = call_api(
                {
                  'method' => 'GET',
                  'uri' => http_client.normalize_uri("/api/v1/namespaces/#{namespace}/pods/#{pod}")
                },
                options
              )

              json
            end

            def list_auth(namespace, options = {})
              data = {
                kind: "SelfSubjectRulesReview",
                "apiVersion": "authorization.k8s.io/v1",
                "metadata": {
                  "creationTimestamp": nil
                },
                "spec": {
                  "namespace": namespace
                },
                "status": {
                  "resourceRules": nil,
                  "nonResourceRules": nil,
                  "incomplete": false
                }
              }

              _res, json = call_api(
                {
                  'method' => 'POST',
                  'uri' => http_client.normalize_uri('/apis/authorization.k8s.io/v1/selfsubjectrulesreviews'),
                  'data' => JSON.pretty_generate(data)
                },
                options
              )

              json
            end

            def get_secret(secret, namespace, options = {})
              _res, json = call_api(
                {
                  'method' => 'GET',
                  'uri' => http_client.normalize_uri("/api/v1/namespaces/#{namespace}/secrets/#{secret}")
                },
                options
              )

              json
            end

            def list_secrets(namespace, options = {})
              _res, json = call_api(
                {
                  'method' => 'GET',
                  'uri' => http_client.normalize_uri("/api/v1/namespaces/#{namespace}/secrets")
                },
                options
              )

              json
            end

            def get_namespace(namespace, options = {})
              _res, json = call_api(
                {
                  'method' => 'GET',
                  'uri' => http_client.normalize_uri("/api/v1/namespaces/#{namespace}")
                },
                options
              )

              json
            end

            def list_namespaces(options = {})
              _res, json = call_api(
                {
                  'method' => 'GET',
                  'uri' => http_client.normalize_uri('/api/v1/namespaces')
                },
                options
              )

              json
            end

            def list_pods(namespace, options = {})
              _res, json = call_api(
                {
                  'method' => 'GET',
                  'uri' => http_client.normalize_uri("/api/v1/namespaces/#{namespace}/pods")
                },
                options
              )

              json
            end

            def create_pod(data, namespace, options = {})
              res, json = call_api(
                {
                  'method' => 'POST',
                  'uri' => http_client.normalize_uri("/api/v1/namespaces/#{namespace}/pods"),
                  'data' => JSON.pretty_generate(data)
                },
                options
              )

              if res.code != 201
                raise Kubernetes::Error::UnexpectedStatusCode, res: res
              end

              json
            end

            def delete_pod(name, namespace, options = {})
              _res, json = call_api(
                {
                  'method' => 'DELETE',
                  'uri' => http_client.normalize_uri("/api/v1/namespaces/#{namespace}/pods/#{name}"),
                  'headers' => {}
                },
                options
              )

              json
            end

            private

            EXEC_CHANNEL_STDIN = 0
            EXEC_CHANNEL_STDOUT = 1
            EXEC_CHANNEL_STDERR = 2
            EXEC_CHANNEL_ERROR = 3
            EXEC_CHANNEL_RESIZE = 4

            attr_reader :http_client

            def call_api(request, options = {})
              res = http_client.send_request_raw(request_options(request, options))

              if res.nil? || res.body.nil?
                raise Kubernetes::Error::InvalidApiError.new(res: res)
              elsif res.code == 401
                raise Kubernetes::Error::AuthenticationError.new(res: res)
              elsif res.code == 403
                raise Kubernetes::Error::ForbiddenError.new(res: res)
              elsif res.code == 404
                raise Kubernetes::Error::NotFoundError.new(res: res)
              elsif res.code >= 500 && res.code <= 599
                raise Kubernetes::Error::ServerError.new(res: res)
              end

              json = res.get_json_document
              if json.nil?
                raise Kubernetes::Error::InvalidApiError.new(res: res)
              end

              [res, json.deep_symbolize_keys]
            end

            def request_options(request, options = {})
              token = options.fetch(:token, @token)

              request.merge(
                {
                  'agent' => USER_AGENT,
                  'headers' => request.fetch('headers', {}).merge(
                    {
                      'Authorization' => "Bearer #{token}"
                    }
                  )
                }
              )
            end
          end
        end
      end
    end
  end
end
