/
gloo.py
101 lines (89 loc) · 3.23 KB
/
gloo.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import traceback
import functools
from urllib.parse import urlparse
import asyncio
import uvloop
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
import re
length_regex = re.compile(r'Content-Length: ([0-9]+)\r\n', re.IGNORECASE)
BUFFER_SIZE = 4096
async def handle_client(reader, writer, loop=None):
header_str = ''
payload = b''
try:
while True:
line = await reader.readline()
if not line or line == b'\r\n':
break
header_str += line.decode('utf-8')
match = length_regex.search(header_str)
if match:
length = int(match.group(1))
while len(payload) < length:
payload += await reader.read(BUFFER_SIZE)
headers = header_str.split('\r\n')[:-1]
method, path, version = headers[0].split(' ')
if method == 'CONNECT': # HTTPS
host, port_str = path.split(':')
port = int(port_str)
rreader, rwriter = await asyncio.open_connection(host=host, port=port, loop=loop)
writer.write(b'HTTP/1.1 200 Connection Established\r\n\r\n')
tasks = [
asyncio.ensure_future(relay(reader, rwriter), loop=loop),
asyncio.ensure_future(relay(rreader, writer), loop=loop)
]
await asyncio.wait(tasks, loop=loop)
else: # HTTP
r = urlparse(path)
host = r.hostname
port = r.port or 80
new_headers = []
new_headers.append(' '.join([method, r.path, version]))
connection_header_found = False
for header in headers[1:]:
name, value = header.strip().split(': ', 1)
if name.lower() == 'proxy-connection':
continue
elif name.lower() == 'connection':
connection_header_found = True
new_headers.append('Connection: close')
else:
new_headers.append(header)
if not connection_header_found:
new_headers.append('Connection: close')
new_headers.append('\r\n')
new_header_str = '\r\n'.join(new_headers)
rreader, rwriter = await asyncio.open_connection(host=host, port=port, loop=loop)
rwriter.write(new_header_str.encode('utf-8'))
if payload:
rwriter.write(payload)
await rwriter.drain()
await relay(rreader, writer)
except:
traceback.print_exc()
finally:
writer.close()
async def relay(reader, writer):
try:
while True:
data = await reader.read(BUFFER_SIZE)
if not data:
break
writer.write(data)
await writer.drain()
except:
traceback.print_exc()
finally:
writer.close()
if __name__ == '__main__':
loop = asyncio.get_event_loop()
try:
callback = functools.partial(handle_client, loop=loop)
loop.run_until_complete(asyncio.start_server(callback, host='127.0.0.1', port=5000, loop=loop))
loop.run_forever()
except KeyboardInterrupt:
print('Closing...')
finally:
loop.close()