/
async.py
102 lines (85 loc) · 2.92 KB
/
async.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
102
import asyncio
import aiohttp
import aiohttp.web as web
import rapidjson as json
from . import Env, log
@asyncio.coroutine
def get_env(request):
env = Env()
resp = yield from aiohttp.request(
'GET', '%s/info/' % env('host_web'),
headers=request.headers,
cookies=request.cookies.items(),
allow_redirects=False
)
if resp.status == 200:
data = yield from resp.json()
username = data.get('username')
if username:
env.username = username
return env, None
else:
return None, web.Response(body=b'403 Forbidden', status=403)
body = yield from resp.read()
return None, web.Response(body=body, status=resp.status)
@asyncio.coroutine
def wshandler(request):
env, error = yield from get_env(request)
if error:
return error
ws = web.WebSocketResponse()
ws.start(request)
request.app['sockets'].append((env.username, ws))
session = request.cookies.get('session')
while True:
msg = yield from ws.receive()
if msg.tp == web.MsgType.text:
log.debug(msg.data)
data = json.loads(msg.data)
payload = data.get('payload')
if payload:
payload = json.dumps(payload)
resp = yield from aiohttp.request(
'POST' if payload else 'GET',
env('host_web') + data['url'],
headers={
'X-Requested-With': 'XMLHttpRequest',
'Cookie': data['cookie']
},
data=payload,
)
log.debug('%s %s', resp.status, msg.data)
if resp.status == 200:
p = (yield from resp.read()).decode()
ws.send_str(json.dumps({'uid': data['uid'], 'payload': p}))
new_session = resp.cookies.get('session')
if new_session and session != new_session:
session = new_session.value
msg = {'session': new_session.output(header='').strip()}
ws.send_str(json.dumps(msg))
log.debug('sent new session')
elif msg.tp == web.MsgType.close:
log.debug('ws closed')
yield from ws.close()
break
elif msg.tp == web.MsgType.error:
log.exception(ws.exception())
request.app['sockets'].remove((env.username, ws))
return ws
@asyncio.coroutine
def notify(request):
env, error = yield from get_env(request)
if error:
return error
yield from request.post()
msg = yield from request.text()
for username, ws in request.app['sockets']:
if username == env.username:
ws.send_str(msg)
return web.Response(body=b'OK')
def create_app():
app = web.Application()
app.router.add_route('GET', '/', wshandler)
app.router.add_route('POST', '/notify/', notify)
app['sockets'] = []
return app