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
103
104
|
#!/usr/bin/env python3
import argparse
import asyncio
import logging
import os
import sys
import time
import websockets.server as ws_server
from websockets.exceptions import ConnectionClosedError
log = logging.getLogger(__name__)
logging.basicConfig(
format="[%(asctime)s] %(message)s",
level=logging.DEBUG,
)
async def echo(websocket):
try:
async for message in websocket:
try:
log.info(f'got request {message}')
except Exception as e:
log.error(f'error {e} getting path from {message}')
await websocket.send(message)
except ConnectionClosedError:
pass
async def on_async_conn(conn):
rpath = str(conn.path)
pcomps = rpath[1:].split('/')
if len(pcomps) == 0:
pcomps = ['echo'] # default handler
log.info(f'connection for {pcomps}')
if pcomps[0] == 'echo':
log.info(f'/echo endpoint')
for message in await conn.recv():
await conn.send(message)
elif pcomps[0] == 'text':
await conn.send('hello!')
elif pcomps[0] == 'file':
if len(pcomps) < 2:
conn.close(code=4999, reason='unknown file')
return
fpath = os.path.join('../', pcomps[1])
if not os.path.exists(fpath):
conn.close(code=4999, reason='file not found')
return
bufsize = 0
if len(pcomps) > 2:
bufsize = int(pcomps[2])
if bufsize <= 0:
bufsize = 16*1024
delay_ms = 0
if len(pcomps) > 3:
delay_ms = int(pcomps[3])
n = 1
if len(pcomps) > 4:
n = int(pcomps[4])
for _ in range(n):
with open(fpath, 'r+b') as fd:
while True:
buf = fd.read(bufsize)
if buf is None or len(buf) == 0:
break
await conn.send(buf)
if delay_ms > 0:
time.sleep(delay_ms/1000)
else:
log.info(f'unknown endpoint: {rpath}')
await conn.close(code=4999, reason='path unknown')
await conn.close(code=1000, reason='')
async def run_server(port):
log.info(f'starting server on port {port}')
async with ws_server.serve(ws_handler=on_async_conn,
host="localhost", port=port):
await asyncio.Future()
async def main():
parser = argparse.ArgumentParser(prog='scorecard',
description="Run a websocket echo server.")
parser.add_argument("--port", type=int,
default=0, help="port to listen on")
args = parser.parse_args()
if args.port == 0:
sys.stderr.write('need --port\n')
sys.exit(1)
logging.basicConfig(
format="%(asctime)s %(message)s",
level=logging.DEBUG,
)
await run_server(args.port)
if __name__ == "__main__":
asyncio.run(main())
|
