class Net::SSH::Service::Forward

def active_local_sockets
  @local_forwarded_sockets.keys
end

def active_locals
  @local_forwarded_ports.keys
end

def active_remote_destinations
  @remote_forwarded_ports.each_with_object({}) do |(remote, local), result|
    result[[local.port, local.host]] = remote
  end
end

def active_remotes
  @remote_forwarded_ports.keys
end

def agent(channel)
  return if @agent_forwarded
  @agent_forwarded = true
  channel.send_channel_request("auth-agent-req@openssh.com") do |achannel, success|
    if success
      debug { "authentication agent forwarding is active" }
    else
      achannel.send_channel_request("auth-agent-req") do |a2channel, success2|
        if success2
          debug { "authentication agent forwarding is active" }
        else
          error { "could not establish forwarding of authentication agent" }
        end
      end
    end
  end
end

def auth_agent_channel(session, channel, packet)
  info { "opening auth-agent channel" }
  channel[:invisible] = true
  begin
    agent = Authentication::Agent.connect(logger, session.options[:agent_socket_factory])
    if (agent.socket.is_a? ::IO)
      prepare_client(agent.socket, channel, :agent)
    else
      prepare_simple_client(agent.socket, channel, :agent)
    end
  rescue Exception => e
    error { "attempted to connect to agent but failed: #{e.class.name} (#{e.message})" }
    raise Net::SSH::ChannelOpenFailed.new(2, "could not connect to authentication agent")
  end
end

def cancel_local(port, bind_address = "127.0.0.1")
  socket = @local_forwarded_ports.delete([port, bind_address])
  socket.shutdown rescue nil
  socket.close rescue nil
  session.stop_listening_to(socket)
end

def cancel_local_socket(local_socket_path)
  socket = @local_forwarded_sockets.delete(local_socket_path)
  socket.shutdown rescue nil
  socket.close rescue nil
  session.stop_listening_to(socket)
end

def cancel_remote(port, host = "127.0.0.1")
  session.send_global_request("cancel-tcpip-forward", :string, host, :long, port) do |success, response|
    if success
      @remote_forwarded_ports.delete([port, host])
    else
      raise Net::SSH::Exception, "could not cancel remote forward request on #{host}:#{port}"
    end
  end
end

def forwarded_tcpip(session, channel, packet)
  connected_address  = packet.read_string
  connected_port     = packet.read_long
  originator_address = packet.read_string
  originator_port    = packet.read_long
  puts "REMOTE 0: #{connected_port} #{connected_address} #{originator_address} #{originator_port}"
  remote = @remote_forwarded_ports[[connected_port, connected_address]]
  if remote.nil?
    raise Net::SSH::ChannelOpenFailed.new(1, "unknown request from remote forwarded connection on #{connected_address}:#{connected_port}")
  end
  puts "REMOTE: #{remote.host} #{remote.port}"
  client = TCPSocket.new(remote.host, remote.port)
  info { "connected #{connected_address}:#{connected_port} originator #{originator_address}:#{originator_port}" }
  prepare_client(client, channel, :remote)
rescue SocketError => err
  raise Net::SSH::ChannelOpenFailed.new(2, "could not connect to remote host (#{remote.host}:#{remote.port}): #{err.message}")
end

def initialize(session)
  @session = session
  self.logger = session.logger
  @remote_forwarded_ports = {}
  @local_forwarded_ports = {}
  @agent_forwarded = false
  @local_forwarded_sockets = {}
  session.on_open_channel('forwarded-tcpip', &method(:forwarded_tcpip))
  session.on_open_channel('auth-agent', &method(:auth_agent_channel))
  session.on_open_channel('auth-agent@openssh.com', &method(:auth_agent_channel))
end

def local(*args)
  if args.length < 3 || args.length > 4
    raise ArgumentError, "expected 3 or 4 parameters, got #{args.length}"
  end
  local_port_type = :long
  socket = begin
    if defined?(UNIXServer) and args.first.class == UNIXServer
      local_port_type = :string
      args.shift
    else
      bind_address = "127.0.0.1"
      bind_address = args.shift if args.first.is_a?(String) && args.first =~ /\D/
      local_port = args.shift.to_i
      local_port_type = :long
      TCPServer.new(bind_address, local_port)
    end
  end
  local_port = socket.addr[1] if local_port == 0 # ephemeral port was requested
  remote_host = args.shift
  remote_port = args.shift.to_i
  @local_forwarded_ports[[local_port, bind_address]] = socket
  session.listen_to(socket) do |server|
    client = server.accept
    debug { "received connection on #{socket}" }
    channel = session.open_channel("direct-tcpip", :string, remote_host, :long,
                                   remote_port, :string, bind_address, local_port_type, local_port) do |achannel|
      achannel.info { "direct channel established" }
    end
    prepare_client(client, channel, :local)
    channel.on_open_failed do |ch, code, description|
      channel.error { "could not establish direct channel: #{description} (#{code})" }
      session.stop_listening_to(channel[:socket])
      channel[:socket].close
    end
  end
  local_port
end

def local_socket(local_socket_path, remote_socket_path)
  File.delete(local_socket_path) if File.exist?(local_socket_path)
  socket = Socket.unix_server_socket(local_socket_path)
  @local_forwarded_sockets[local_socket_path] = socket
  session.listen_to(socket) do |server|
    client = server.accept[0]
    debug { "received connection on #{socket}" }
    channel = session.open_channel("direct-streamlocal@openssh.com",
                                   :string, remote_socket_path,
                                   :string, nil,
                                   :long, 0) do |achannel|
      achannel.info { "direct channel established" }
    end
    prepare_client(client, channel, :local)
    channel.on_open_failed do |ch, code, description|
      channel.error { "could not establish direct channel: #{description} (#{code})" }
      session.stop_listening_to(channel[:socket])
      channel[:socket].close
    end
  end
  local_socket_path
end

def prepare_client(client, channel, type)
  client.extend(Net::SSH::BufferedIo)
  client.extend(Net::SSH::ForwardedBufferedIo)
  client.logger = logger
  session.listen_to(client)
  channel[:socket] = client
  channel.on_data do |ch, data|
    debug { "data:#{data.length} on #{type} forwarded channel" }
    ch[:socket].enqueue(data)
  end
  channel.on_eof do |ch|
    debug { "eof #{type} on #{type} forwarded channel" }
    begin
      ch[:socket].send_pending
      ch[:socket].shutdown Socket::SHUT_WR
    rescue IOError => e
      if e.message =~ /closed/ then
        debug { "epipe in on_eof => shallowing exception:#{e}" }
      else
        raise
      end
    rescue Errno::EPIPE => e
      debug { "epipe in on_eof => shallowing exception:#{e}" }
    rescue Errno::ENOTCONN => e
      debug { "enotconn in on_eof => shallowing exception:#{e}" }
    end
  end
  channel.on_close do |ch|
    debug { "closing #{type} forwarded channel" }
    ch[:socket].close if !client.closed?
    session.stop_listening_to(ch[:socket])
  end
  channel.on_process do |ch|
    if ch[:socket].closed?
      ch.info { "#{type} forwarded connection closed" }
      ch.close
    elsif ch[:socket].available > 0
      data = ch[:socket].read_available(8192)
      ch.debug { "read #{data.length} bytes from client, sending over #{type} forwarded connection" }
      ch.send_data(data)
    end
  end
end

def prepare_simple_client(client, channel, type)
  channel[:socket] = client
  channel.on_data do |ch, data|
    ch.debug { "data:#{data.length} on #{type} forwarded channel" }
    ch[:socket].send(data)
  end
  channel.on_process do |ch|
    data = ch[:socket].read(8192)
    if data
      ch.debug { "read #{data.length} bytes from client, sending over #{type} forwarded connection" }
      ch.send_data(data)
    end
  end
end

def remote(port, host, remote_port, remote_host = "127.0.0.1")
  session.send_global_request("tcpip-forward", :string, remote_host, :long, remote_port) do |success, response|
    if success
      remote_port = response.read_long if remote_port == 0
      debug { "remote forward from remote #{remote_host}:#{remote_port} to #{host}:#{port} established" }
      @remote_forwarded_ports[[remote_port, remote_host]] = Remote.new(host, port)
      yield remote_port, remote_host if block_given?
    else
      instruction = if block_given?
                      yield :error
                    end
      unless instruction == :no_exception
        error { "remote forwarding request failed" }
        raise Net::SSH::Exception, "remote forwarding request failed"
      end
    end
  end
end