#!/usr/bin/env python
# SOCKSv5 proxy server for network-sandbox
# Copyright 2015 Gentoo Foundation
# Distributed under the terms of the GNU General Public License v2

import asyncio
import errno
import os
import socket
import struct
import sys

if hasattr(asyncio, 'ensure_future'):
	# Python >=3.4.4.
	asyncio_ensure_future = asyncio.ensure_future
else:
	# getattr() necessary because async is a keyword in Python >=3.7.
	asyncio_ensure_future = getattr(asyncio, 'async')

try:
	current_task = asyncio.current_task
except AttributeError:
	# Deprecated since Python 3.7
	current_task = asyncio.Task.current_task


class Socks5Server:
	"""
	An asynchronous SOCKSv5 server.
	"""

	async def handle_proxy_conn(self, reader, writer):
		"""
		Handle incoming client connection. Perform SOCKSv5 request
		exchange, open a proxied connection and start relaying.

		@param reader: Read side of the socket
		@type reader: asyncio.StreamReader
		@param writer: Write side of the socket
		@type writer: asyncio.StreamWriter
		"""

		try:
			# SOCKS hello
			data = await reader.readexactly(2)
			vers, method_no = struct.unpack('!BB', data)

			if vers != 0x05:
				# disconnect on invalid packet -- we have no clue how
				# to reply in alien :)
				writer.close()
				return

			# ...and auth method list
			data = await reader.readexactly(method_no)
			for method in data:
				if method == 0x00:
					break
			else:
				# no supported method
				method = 0xFF

			# auth reply
			repl = struct.pack('!BB', 0x05, method)
			writer.write(repl)
			await writer.drain()
			if method == 0xFF:
				writer.close()
				return

			# request
			data = await reader.readexactly(4)
			vers, cmd, rsv, atyp = struct.unpack('!BBBB', data)

			if vers != 0x05 or rsv != 0x00:
				# disconnect on malformed packet
				self.close()
				return

			# figure out if we can handle it
			rpl = 0x00
			if cmd != 0x01:  # CONNECT
				rpl = 0x07  # command not supported
			elif atyp == 0x01:  # IPv4
				data = await reader.readexactly(4)
				addr = socket.inet_ntoa(data)
			elif atyp == 0x03:  # domain name
				data = await reader.readexactly(1)
				addr_len, = struct.unpack('!B', data)
				addr = await reader.readexactly(addr_len)
				try:
					addr = addr.decode('idna')
				except UnicodeDecodeError:
					rpl = 0x04  # host unreachable

			elif atyp == 0x04:  # IPv6
				data = await reader.readexactly(16)
				addr = socket.inet_ntop(socket.AF_INET6, data)
			else:
				rpl = 0x08  # address type not supported

			# try to connect if we can handle it
			if rpl == 0x00:
				data = await reader.readexactly(2)
				port, = struct.unpack('!H', data)

				try:
					# open a proxied connection
					proxied_reader, proxied_writer = await asyncio.open_connection(
							addr, port)
				except (socket.gaierror, socket.herror):
					# DNS failure
					rpl = 0x04  # host unreachable
				except OSError as e:
					# connection failure
					if e.errno in (errno.ENETUNREACH, errno.ENETDOWN):
						rpl = 0x03  # network unreachable
					elif e.errno in (errno.EHOSTUNREACH, errno.EHOSTDOWN):
						rpl = 0x04  # host unreachable
					elif e.errno in (errno.ECONNREFUSED, errno.ETIMEDOUT):
						rpl = 0x05  # connection refused
					else:
						raise
				else:
					# get socket details that we can send back to the client
					# local address (sockname) in particular -- but we need
					# to ask for the whole socket since Python's sockaddr
					# does not list the family...
					sock = proxied_writer.get_extra_info('socket')
					addr = sock.getsockname()
					if sock.family == socket.AF_INET:
						host, port = addr
						bin_host = socket.inet_aton(host)

						repl_addr = struct.pack('!B4sH',
								0x01, bin_host, port)
					elif sock.family == socket.AF_INET6:
						# discard flowinfo, scope_id
						host, port = addr[:2]
						bin_host = socket.inet_pton(sock.family, host)

						repl_addr = struct.pack('!B16sH',
								0x04, bin_host, port)

			if rpl != 0x00:
				# fallback to 0.0.0.0:0
				repl_addr = struct.pack('!BLH', 0x01, 0x00000000, 0x0000)

			# reply to the request
			repl = struct.pack('!BBB', 0x05, rpl, 0x00)
			writer.write(repl + repl_addr)
			await writer.drain()

			# close if an error occured
			if rpl != 0x00:
				writer.close()
				return

			# otherwise, start two loops:
			# remote -> local...
			t = asyncio_ensure_future(self.handle_proxied_conn(
					proxied_reader, writer, current_task()))

			# and local -> remote...
			try:
				try:
					while True:
						data = await reader.read(4096)
						if data == b'':
							# client disconnected, stop relaying from
							# remote host
							t.cancel()
							break

						proxied_writer.write(data)
						await proxied_writer.drain()
				except OSError:
					# read or write failure
					t.cancel()
				except:
					t.cancel()
					raise
			finally:
				# always disconnect in the end :)
				proxied_writer.close()
				writer.close()

		except (OSError, asyncio.IncompleteReadError, asyncio.CancelledError):
			writer.close()
			return
		except:
			writer.close()
			raise

	async def handle_proxied_conn(self, proxied_reader, writer, parent_task):
		"""
		Handle the proxied connection. Relay incoming data
		to the client.

		@param reader: Read side of the socket
		@type reader: asyncio.StreamReader
		@param writer: Write side of the socket
		@type writer: asyncio.StreamWriter
		"""

		try:
			try:
				while True:
					data = await proxied_reader.read(4096)
					if data == b'':
						break

					writer.write(data)
					await writer.drain()
			finally:
				parent_task.cancel()
		except (OSError, asyncio.CancelledError):
			return


if __name__ == '__main__':
	if len(sys.argv) != 2:
		print('Usage: %s <socket-path>' % sys.argv[0])
		sys.exit(1)

	loop = asyncio.get_event_loop()
	s = Socks5Server()
	server = loop.run_until_complete(
		asyncio.start_unix_server(s.handle_proxy_conn, sys.argv[1]))

	ret = 0
	try:
		try:
			loop.run_forever()
		except KeyboardInterrupt:
			pass
		except:
			ret = 1
	finally:
		server.close()
		loop.run_until_complete(server.wait_closed())
		loop.close()
		os.unlink(sys.argv[1])