# frozen_string_literal: true

# vim: set et sw=2:

require "net/http"
require "uri"
require "socket"

require_relative "rsocks5patch/version"

# When this module is evaluated it will patch the Ruby Socket class.
module Rsocks5patch
  class Error < StandardError; end

  Socket.class_eval do
    # alias :orig_func :self.tcp
    #
    class << self
      alias_method :orig_func, :tcp

      attr_accessor :socks_host, :socks_port
    end

    def self.mustread(socket, num)
      data = ""
      loop do
        tmp = socket.read(num - data.length)
        return null if tmp.nil?

        data += tmp
        return data if data.length == num
      end
    end

    def self.tcp(
      host,
      port,
      local_host = nil,
      local_port = nil,
      connect_timeout: nil,
      resolv_timeout: nil,
      socks_host: nil,
      socks_port: nil
    )
      socks_host ||= self.socks_host
      socks_port ||= self.socks_port

      unless socks_host
        return orig_func(host, port, local_host, local_port, connect_timeout: connect_timeout,
                                                             resolv_timeout: resolv_timeout)
      end

      socket = orig_func(socks_host, socks_port, local_host, local_port, connect_timeout: nil, resolv_timeout: nil)

      socket.write([0x5, 0x1, 0x0].pack("CCC"))
      ver, auth = mustread(socket, 2).unpack("CC")

      raise StandardError, "Unexpected version #{ver}." if ver != 0x5
      raise StandardError, "No available auth method." if auth == 0xff

      # Version, CMD=1, RSV=0, ADDRTYPE=3, LEN, HOST, PORT
      socket.write([0x5, 0x01, 0x00, 0x3, host.length, host, port].pack("CCCCCa*n"))

      ver, status, _resv, addrtype = mustread(socket, 4).unpack("CCCC")
      raise StandardError, "Unexpected version #{ver}." if ver != 0x5

      case status
      when 0x00
        ""
      when 0x01
        raise StandardError, "general failure"
      when 0x02
        raise StandardError, "connection not allowed by ruleset"
      when 0x03
        raise StandardError, "network unreachable"
      when 0x04
        raise StandardError, "host unreachable"
      when 0x05
        raise StandardError, "connection refused by destination host"
      when 0x06
        raise StandardError, "TTL expired"
      when 0x07
        raise StandardError, "command not supported / protocol error"
      when 0x08
        raise StandardError, "address type not supported"
      else
        raise StandardError, "unexpected error status #{status}"
      end

      case addrtype
      when 0x1
        _ipv4addr = mustread(socket, 4)
      when 0x3
        _len = mustread(socket, 1)
        _host = mustread(socket, len)
      when 0x4
        _ipv6addr = mustread(socket, 16)
      else
        raise StandardError, "No available auth method." if auth == 0xff
      end
      _port = mustread(socket, 2)

      socket
    end
  end
end