-
Notifications
You must be signed in to change notification settings - Fork 7
Expand file tree
/
Copy pathnrepl_bencode.py
More file actions
108 lines (97 loc) · 3.34 KB
/
Copy pathnrepl_bencode.py
File metadata and controls
108 lines (97 loc) · 3.34 KB
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
105
106
107
108
#!/usr/bin/env python3
"""Minimal nREPL bencode client, shared by the completion oracle harness
(scripts/completion_oracle.py) and ad-hoc probing.
The encoder/decoder mirrors the inline implementations the nREPL e2e
scripts carry (test/e2e/phase14_nrepl.sh) — kept dependency-free so any
python3 can run it. Responses are plain dicts/lists/strs/ints; byte
strings are decoded as UTF-8 (replacement on error).
"""
import socket
def bencode(x):
if isinstance(x, dict):
return b"d" + b"".join(bencode(k) + bencode(v) for k, v in sorted(x.items())) + b"e"
if isinstance(x, str):
b = x.encode()
return str(len(b)).encode() + b":" + b
if isinstance(x, int):
return b"i" + str(x).encode() + b"e"
if isinstance(x, list):
return b"l" + b"".join(bencode(v) for v in x) + b"e"
raise TypeError(f"bencode: {type(x)}")
def bdecode(buf, i=0):
"""Decode one value at offset i. Returns (value, next_offset).
Raises ValueError/IndexError on a truncated buffer."""
c = buf[i:i + 1]
if c == b"d":
d = {}
i += 1
while buf[i:i + 1] != b"e":
k, i = bdecode(buf, i)
v, i = bdecode(buf, i)
d[k] = v
return d, i + 1
if c == b"l":
out = []
i += 1
while buf[i:i + 1] != b"e":
v, i = bdecode(buf, i)
out.append(v)
return out, i + 1
if c == b"i":
j = buf.index(b"e", i)
return int(buf[i + 1:j]), j + 1
if c == b"":
raise ValueError("empty")
j = buf.index(b":", i)
n = int(buf[i:j])
end = j + 1 + n
if end > len(buf):
raise ValueError("truncated string")
return buf[j + 1:end].decode("utf-8", "replace"), end
class NreplClient:
"""One nREPL connection; request() sends an op and collects response
dicts until a `done` status arrives (or the socket times out)."""
def __init__(self, host, port, timeout=20):
self.sock = socket.create_connection((host, port), timeout=10)
self.sock.settimeout(timeout)
self.buf = b""
self.counter = 0
self.session = None
def request(self, msg):
self.counter += 1
msg = dict(msg)
msg.setdefault("id", f"req{self.counter}")
if self.session and "session" not in msg:
msg["session"] = self.session
self.sock.sendall(bencode(msg))
out = []
while True:
try:
chunk = self.sock.recv(1 << 16)
except socket.timeout:
break
if not chunk:
break
self.buf += chunk
progressed = True
while progressed:
progressed = False
try:
v, n = bdecode(self.buf)
except (ValueError, IndexError):
continue
out.append(v)
self.buf = self.buf[n:]
progressed = True
if any(isinstance(r, dict) and "done" in r.get("status", []) for r in out):
break
return out
def clone(self):
res = self.request({"op": "clone"})
self.session = next(r["new-session"] for r in res if isinstance(r, dict) and "new-session" in r)
return self.session
def close(self):
try:
self.sock.close()
except OSError:
pass