Created
September 5, 2017 01:02
-
-
Save Sherlock-Holo/468b85feca5599f9baab0df65c64cd13 to your computer and use it in GitHub Desktop.
asyncio can't use uvloop
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python3 | |
import asyncio | |
import functools | |
import logging | |
import socket | |
import struct | |
logging.basicConfig( | |
level=logging.DEBUG, | |
format='{asctime} {levelname} {message}', | |
datefmt='%Y-%m-%d %H:%M:%S', | |
style='{') | |
class Server: | |
async def handle(self, reader, writer): | |
logging.debug('connect from {}'.format(writer.get_extra_info('peername'))) | |
request = await reader.read(2) | |
if request[0] != 5: | |
writer.close() | |
logging.error('socks version not support') | |
return None | |
else: | |
nmethods = request[1] | |
logging.debug('methods number: {}'.format(nmethods)) | |
methods = await reader.read(nmethods) | |
if 0 in methods: | |
writer.write(b'\x05\x00') | |
await writer.drain() | |
else: | |
writer.write(b'\x05\xff') | |
logging.error('Authentication not support') | |
writer.close() | |
return None | |
data = await reader.read(4) | |
ver, cmd, rsv, atyp = data | |
if cmd != 1: | |
data = [] | |
data.append(b'\x05\x07\x00\x01') | |
data.append(socket.inet_aton('0.0.0.0')) | |
data.append(struct.pack('>H', 0)) | |
writer.write(b''.join(data)) | |
writer.close() | |
logging.error('cmd not support') | |
return None | |
else: | |
if atyp == 1: | |
_addr = await reader.read(4) | |
addr = socket.inet_ntoa(_addr) | |
elif atyp == 3: | |
addr_len = await reader.read(1) | |
addr = await reader.read(ord(addr_len)) | |
elif atyp == 4: | |
_addr = await reader.read(16) | |
addr = socket.inet_ntop(socket.AF_INET6, _addr) | |
_port = await reader.read(2) | |
port = struct.unpack('>H', _port)[0] | |
logging.debug('remote: {}:{}'.format(addr, port)) | |
try: | |
r_reader, r_writer = await asyncio.open_connection(addr, port) | |
except: | |
data = [] | |
data.append(b'\x05\x04\x00\x01') | |
data.append(socket.inet_aton('0.0.0.0')) | |
data.append(struct.pack('>H', 0)) | |
writer.write(b''.join(data)) | |
writer.close() | |
return None | |
bind_info = r_writer.get_extra_info('sockname') | |
data = [] | |
data.append(b'\x05\x00\x00') | |
try: | |
bind_addr = socket.inet_pton(socket.AF_INET6, bind_info[0]) | |
data.append(b'\x04') | |
except OSError: | |
bind_addr = socket.inet_aton(bind_info[0]) | |
data.append(b'\x01') | |
data.append(bind_addr) | |
data.append(struct.pack('>H', bind_info[1])) | |
writer.write(b''.join(data)) | |
await writer.drain() | |
async def relay(fr, to): | |
try: | |
while True: | |
data = await fr.read(4096) | |
if not data: | |
break | |
to.write(data) | |
await to.drain() | |
except OSError as e: | |
logging.error(e) | |
return None | |
logging.debug('start relay') | |
s2r = asyncio.ensure_future(relay(reader, r_writer)) | |
r2s = asyncio.ensure_future(relay(r_reader, writer)) | |
s2r.add_done_callback(functools.partial(self.close_transport, writer)) | |
r2s.add_done_callback(functools.partial(self.close_transport, r_writer)) | |
def close_transport(self, writer, future): | |
writer.close() | |
if __name__ == '__main__': | |
s = Server() | |
try: | |
import uvloop | |
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) | |
except ImportError: | |
pass | |
loop = asyncio.get_event_loop() | |
coro = asyncio.start_server(s.handle, '127.0.0.2', 1089, loop=loop) | |
server = loop.run_until_complete(coro) | |
try: | |
loop.run_forever() | |
except KeyboardInterrupt: | |
pass | |
server.close() | |
loop.run_until_complete(server.wait_closed()) | |
loop.close() |
Author
Sherlock-Holo
commented
Sep 5, 2017
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment