Skip to content

Commit e32fc8a

Browse files
isabelatkinsonisabelatkinson
committed
move script
1 parent 1921fef commit e32fc8a

File tree

2 files changed

+280
-2
lines changed

2 files changed

+280
-2
lines changed

‎.evergreen/run-socks5-proxy-test.sh‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,10 @@ FEATURE_FLAGS+=("socks5-proxy")
99
CARGO_OPTIONS+=("--ignore-default-filter")
1010

1111
# with auth
12-
python $DRIVERS_TOOLS/.evergreen/socks5srv.py --map "localhost:12345 to localhost:27017" --port 1080 --auth username:p4ssw0rd &
12+
python .evergreen/socks5srv.py --map "localhost:12345 to localhost:27017" --port 1080 --auth username:p4ssw0rd &
1313
AUTH_PID=$!
1414
# without auth
15-
python $DRIVERS_TOOLS/.evergreen/socks5srv.py --map "localhost:12345 to localhost:27017" --port 1081 &
15+
python .evergreen/socks5srv.py --map "localhost:12345 to localhost:27017" --port 1081 &
1616
NOAUTH_PID=$!
1717

1818
echo "testing socks5proxy with URI $MONGODB_URI"

‎.evergreen/socks5srv.py‎

Lines changed: 278 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,278 @@
1+
#!/usr/bin/env python3
2+
import argparse
3+
import re
4+
import select
5+
import socket
6+
import socketserver
7+
8+
# Usage: python3 socks5srv.py --port port [--auth username:password] [--map 'host:port to host:port' ...]
9+
10+
11+
class AddressRemapper:
12+
"""A helper for remapping (host, port) tuples to new (host, port) tuples
13+
14+
This is useful for Socks5 servers used in testing environments,
15+
because the successful use of the Socks5 proxy can be demonstrated
16+
by being able to 'connect' to a redirected port, which would always
17+
fail without the proxy, even on localhost-only environments
18+
"""
19+
20+
def __init__(self, mappings):
21+
self.mappings = [
22+
AddressRemapper.parse_single_mapping(string) for string in mappings
23+
]
24+
self.add_dns_remappings()
25+
26+
@staticmethod
27+
def parse_single_mapping(string):
28+
print(string)
29+
"""Parse a single mapping of the for '{host}:{port} to {host}:{port}'"""
30+
31+
# Accept either [ipv6]:port or host:port
32+
host_re = r"(\[(?P<{0}_ipv6>[^[\]]+)\]|(?P<{0}_host>[^\[]+))"
33+
port_re = r"(?P<{0}_port>\d+)"
34+
35+
src_re = host_re.format("src") + ":" + port_re.format("src")
36+
dst_re = host_re.format("dst") + ":" + port_re.format("dst")
37+
full_re = "^" + src_re + " to " + dst_re + "$"
38+
39+
match = re.match(full_re, string)
40+
if match is None:
41+
raise Exception(
42+
f"Mapping {string} does not match format '{{host}}:{{port}} to {{host}}:{{port}}'"
43+
)
44+
45+
src = (
46+
(match.group("src_ipv6") or match.group("src_host")).encode("utf8"),
47+
int(match.group("src_port")),
48+
)
49+
dst = (
50+
(match.group("dst_ipv6") or match.group("dst_host")).encode("utf8"),
51+
int(match.group("dst_port")),
52+
)
53+
return (src, dst)
54+
55+
def add_dns_remappings(self):
56+
"""Add mappings for the IP addresses corresponding to hostnames
57+
58+
For example, if there is a mapping (localhost, 1000) to (localhost, 2000),
59+
then this also adds (127.0.0.1, 1000) to (localhost, 2000)."""
60+
61+
for src, dst in self.mappings:
62+
host, port = src
63+
try:
64+
addrs = socket.getaddrinfo(
65+
host, port, socket.AF_UNSPEC, socket.SOCK_STREAM
66+
)
67+
except socket.gaierror:
68+
continue
69+
70+
existing_src_entries = [src for src, dst in self.mappings]
71+
for af, socktype, proto, canonname, sa in addrs:
72+
if af == socket.AF_INET and sa not in existing_src_entries:
73+
self.mappings.append((sa, dst))
74+
elif af == socket.AF_INET6 and sa[:2] not in existing_src_entries:
75+
self.mappings.append((sa[:2], dst))
76+
77+
def remap(self, hostport):
78+
"""Re-map a (host, port) tuple to a new (host, port) tuple if that was requested"""
79+
80+
for src, dst in self.mappings:
81+
if hostport == src:
82+
return dst
83+
return hostport
84+
85+
86+
class Socks5Server(socketserver.ThreadingTCPServer):
87+
"""A simple Socks5 proxy server"""
88+
89+
def __init__(self, server_address, RequestHandlerClass, args):
90+
socketserver.ThreadingTCPServer.__init__(
91+
self, server_address, RequestHandlerClass
92+
)
93+
self.args = args
94+
self.address_remapper = AddressRemapper(args.map)
95+
96+
97+
class Socks5Handler(socketserver.BaseRequestHandler):
98+
"""Request handler for Socks5 connections"""
99+
100+
def finish(self):
101+
"""Called after handle(), always just closes the connection"""
102+
103+
self.request.close()
104+
105+
def read_exact(self, n):
106+
"""Read n bytes from a socket
107+
108+
In Socks5, strings are prefixed with a single byte containing
109+
their length. This method reads a bytes string containing n bytes
110+
(where n can be a number or a bytes object containing that
111+
single byte).
112+
113+
If reading from the client ends prematurely, this returns None.
114+
"""
115+
116+
if type(n) is bytes:
117+
if len(n) == 0:
118+
return None
119+
assert len(n) == 1
120+
n = n[0]
121+
122+
buf = bytearray(n)
123+
mv = memoryview(buf)
124+
bytes_read = 0
125+
while bytes_read < n:
126+
try:
127+
chunk_length = self.request.recv_into(mv[bytes_read:])
128+
except OSError:
129+
return None
130+
if chunk_length == 0:
131+
return None
132+
133+
bytes_read += chunk_length
134+
return bytes(buf)
135+
136+
def create_outgoing_tcp_connection(self, dst, port):
137+
"""Create an outgoing TCP connection to dst:port"""
138+
139+
outgoing = None
140+
for res in socket.getaddrinfo(dst, port, socket.AF_UNSPEC, socket.SOCK_STREAM):
141+
af, socktype, proto, canonname, sa = res
142+
try:
143+
outgoing = socket.socket(af, socktype, proto)
144+
except OSError:
145+
continue
146+
try:
147+
outgoing.connect(sa)
148+
except OSError:
149+
outgoing.close()
150+
continue
151+
break
152+
return outgoing
153+
154+
def handle(self):
155+
"""Handle the Socks5 communication with a freshly connected client"""
156+
157+
# This implements the Socks5 protocol as specified in
158+
# https://datatracker.ietf.org/doc/html/rfc1928
159+
# and username/password authentication as specified in
160+
# https://datatracker.ietf.org/doc/html/rfc1929
161+
# If you prefer HTML tables over ASCII tables, Wikipedia
162+
# also currently has a decent description of the protocol in
163+
# https://en.wikipedia.org/wiki/SOCKS#SOCKS5.
164+
165+
# Receive/send errors are intentionally left unhandled. Closing
166+
# the socket is just fine in that case for us.
167+
168+
# Client greeting
169+
if self.request.recv(1) != b"\x05": # Socks5 only
170+
return
171+
n_auth = self.request.recv(1)
172+
client_auth_methods = self.read_exact(n_auth)
173+
if client_auth_methods is None:
174+
return
175+
176+
# choose either no-auth or username/password
177+
required_auth_method = b"\x00" if self.server.args.auth is None else b"\x02"
178+
if required_auth_method not in client_auth_methods:
179+
self.request.sendall(b"\x05\xff")
180+
return
181+
182+
self.request.sendall(b"\x05" + required_auth_method)
183+
if required_auth_method == b"\x02":
184+
auth_version = self.request.recv(1)
185+
if auth_version != b"\x01": # Only username/password auth v1
186+
return
187+
username_len = self.request.recv(1)
188+
username = self.read_exact(username_len)
189+
password_len = self.request.recv(1)
190+
password = self.read_exact(password_len)
191+
if username is None or password is None:
192+
return
193+
if (
194+
username.decode("utf8") + ":" + password.decode("utf8")
195+
!= self.server.args.auth
196+
):
197+
return
198+
self.request.sendall(b"\x01\x00") # auth success
199+
200+
if self.request.recv(1) != b"\x05": # Socks5 only
201+
return
202+
if self.request.recv(1) != b"\x01": # Outgoing TCP only
203+
return
204+
if self.request.recv(1) != b"\x00": # Reserved, must be 0
205+
return
206+
207+
addrtype = self.request.recv(1)
208+
dst = None
209+
if addrtype == b"\x01": # IPv4
210+
ipv4raw = self.read_exact(4)
211+
if ipv4raw is not None:
212+
dst = ".".join(["{}"] * 4).format(*ipv4raw)
213+
elif addrtype == b"\x03": # Domain
214+
domain_len = self.request.recv(1)
215+
dst = self.read_exact(domain_len)
216+
elif addrtype == b"\x04": # IPv6
217+
ipv6raw = self.read_exact(16)
218+
if ipv6raw is not None:
219+
dst = ":".join(["{:0>2x}{:0>2x}"] * 8).format(*ipv6raw)
220+
else:
221+
return
222+
223+
if dst is None:
224+
return
225+
226+
portraw = self.read_exact(2)
227+
port = portraw[0] * 256 + portraw[1]
228+
229+
(dst, port) = self.server.address_remapper.remap((dst, port))
230+
231+
outgoing = self.create_outgoing_tcp_connection(dst, port)
232+
if outgoing is None:
233+
self.request.sendall(b"\x05\x01\x00") # just report a general failure
234+
return
235+
# success response, do not bother actually stating the locally bound
236+
# host/port address and instead always say 127.0.0.1:4096.
237+
# for our use case, the client will not be making meaningful use
238+
# of this anyway
239+
self.request.sendall(b"\x05\x00\x00\x01\x7f\x00\x00\x01\x10\x00")
240+
241+
self.raw_proxy(self.request, outgoing)
242+
243+
def raw_proxy(self, a, b):
244+
"""Proxy data between sockets a and b as-is"""
245+
246+
with a, b:
247+
while True:
248+
try:
249+
(readable, _, _) = select.select([a, b], [], [])
250+
except (OSError, ValueError):
251+
return
252+
253+
if not readable:
254+
continue
255+
for sock in readable:
256+
buf = sock.recv(4096)
257+
if buf == b"":
258+
return
259+
if sock is a:
260+
b.sendall(buf)
261+
else:
262+
a.sendall(buf)
263+
264+
265+
def main():
266+
parser = argparse.ArgumentParser(description="Start a Socks5 proxy server.")
267+
parser.add_argument("--port", type=int, required=True)
268+
parser.add_argument("--auth", type=str)
269+
parser.add_argument("--map", type=str, action="append", default=[])
270+
args = parser.parse_args()
271+
272+
socketserver.TCPServer.allow_reuse_address = True
273+
with Socks5Server(("localhost", args.port), Socks5Handler, args) as server:
274+
server.serve_forever()
275+
276+
277+
if __name__ == "__main__":
278+
main()

0 commit comments

Comments
 (0)