require 'evil-proxy/httpproxy'
require 'evil-proxy/agentproxy'
require 'evil-proxy/quickcert'

class EvilProxy::MITMProxyServer < EvilProxy::HTTPProxyServer

  def initialize config
    super
    @mitm_pattern = config[:MITMPattern]
    @mitm_servers = {}
    @mitm_port = 4433
  end

  def ca
    return @ca if @ca
    logger.info "Create CA root cert"

    ca_config = {}
    ca_config[:hostname] = 'ca'
    ca_config[:domainname] = 'mitm.proxy'
    ca_config[:password] = 'password'
    ca_config[:CA_dir] ||= File.join(Dir.pwd, "certs/CA")

    ca_config[:keypair_file] ||= File.join ca_config[:CA_dir], "private/cakeypair.pem"
    ca_config[:cert_file] ||= File.join ca_config[:CA_dir], "cacert.pem"
    ca_config[:serial_file] ||= File.join ca_config[:CA_dir], "serial"
    ca_config[:new_certs_dir] ||= File.join ca_config[:CA_dir], "newcerts"
    ca_config[:new_keypair_dir] ||= File.join ca_config[:CA_dir], "private/keypair_backup"
    ca_config[:crl_dir] ||= File.join ca_config[:CA_dir], "crl"

    ca_config[:ca_cert_days] ||= 5 * 365 # five years
    ca_config[:ca_rsa_key_length] ||= 2048

    ca_config[:cert_days] ||= 365 # one year
    ca_config[:cert_key_length_min] ||= 1024
    ca_config[:cert_key_length_max] ||= 2048

    ca_config[:crl_file] ||= File.join ca_config[:crl_dir], "#{ca_config[:hostname]}.crl"
    ca_config[:crl_pem_file] ||= File.join ca_config[:crl_dir], "#{ca_config[:hostname]}.pem"
    ca_config[:crl_days] ||= 14

    if ca_config[:name].nil?
      ca_config[:name] = [
        ['C', 'US', OpenSSL::ASN1::PRINTABLESTRING],
        ['O', ca_config[:domainname], OpenSSL::ASN1::UTF8STRING],
        ['OU', ca_config[:hostname], OpenSSL::ASN1::UTF8STRING],
      ]
    end

    @ca = QuickCert.new ca_config
  end

  def create_self_signed_cert host
    cn = [["C", "US"], ["O", host], ["CN", host]]
    comment = "Generated by Ruby/OpenSSL/MITMProxyServer"
    name = OpenSSL::X509::Name.new(cn)
    hostname = name.to_s.scan(/CN=([\w.]+)/)[0][0]

    logger.info "Create cert for #{hostname}"
    cert_config = { type: 'server', hostname: hostname }
    cert_file, cert, key = ca.create_cert(cert_config)

    return cert, key
  end

  def retry_start_agent_server config
    mitm_server = nil
    10.times do
      begin
        # XXX: ask system for an unused port
        config = config.merge(Port: @mitm_port)
        mitm_server = EvilProxy::AgentProxyServer.new config
      rescue Errno::EADDRINUSE
      rescue Errno::EINVAL => e
        logger.error e.message
        return
      ensure
        @mitm_port += 1
        return mitm_server if mitm_server
      end
    end
    raise RuntimeError, "No avaliable port found, stop retrying"
  end

  def start_mitm_server host, port
    if @mitm_servers[host]
      return @mitm_servers[host].config[:Port]
    else
      cert, key = create_self_signed_cert host
      agent_config = self.config.merge(
        MITMProxyServer: self,
        SSLEnable: true,
        SSLVerifyClient: OpenSSL::SSL::VERIFY_NONE,
        SSLCertificate: cert,
        SSLPrivateKey: key,
      )
      mitm_server = retry_start_agent_server agent_config

      @mitm_servers[host] = mitm_server

      Thread.new do mitm_server.start end
      return mitm_server.config[:Port]
    end
  end

  define_callback_methods :before_mitm
  define_callback_methods :after_mitm

  def do_MITM req, res
    fire :before_mitm, req

    host, port = req.unparsed_uri.split(":")
    port ||= 443

    mitm_port = start_mitm_server host, port
    req.unparsed_uri = "127.0.0.1:#{mitm_port}"

    fire :after_mitm, req, res
  end

  def do_CONNECT req, res
    if !@mitm_pattern || req.unparsed_uri =~ @mitm_pattern
      do_MITM req, res
    end
    super
  end

end
