|
1 | 1 | from blpapi_import_helper import blpapi |
2 | 2 | from argparse import Action |
| 3 | +from collections import namedtuple |
3 | 4 |
|
4 | 5 |
|
5 | 6 | _SESSION_IDENTITY_AUTH_OPTIONS = "sessionIdentityAuthOptions" |
@@ -69,16 +70,38 @@ def __call__(self, parser, args, values, option_string=None): |
69 | 70 | setattr(args, _SESSION_IDENTITY_AUTH_OPTIONS, authOptions) |
70 | 71 |
|
71 | 72 |
|
| 73 | +HostPort = namedtuple("HostPort", ["host", "port"]) |
| 74 | +ServerAddress = namedtuple("ServerAddress", ["endpoint", "socks5"]) |
| 75 | + |
| 76 | + |
72 | 77 | class HostAction(Action): |
73 | 78 | """The action that parses host options from user input""" |
74 | 79 |
|
75 | | - def __call__(self, parser, args, values, option_string=None): |
76 | | - vals = values.split(":", 1) |
| 80 | + @staticmethod |
| 81 | + def _parseHostPort(parser, value): |
| 82 | + vals = value.split(":", 1) |
77 | 83 | if len(vals) != 2: |
| 84 | + parser.error(f"Invalid host option '{value}'") |
| 85 | + |
| 86 | + maxPortValue = 65535 |
| 87 | + if int(vals[1]) <= 0 or int(vals[1]) > maxPortValue: |
| 88 | + parser.error( |
| 89 | + f"Invalid port '{value}', value must be 1 through {maxPortValue}" |
| 90 | + ) |
| 91 | + return HostPort(vals[0], vals[1]) |
| 92 | + |
| 93 | + def __call__(self, parser, args, values, option_string=None): |
| 94 | + endpointSocks5 = values.split("/", 1) |
| 95 | + if len(endpointSocks5) == 0 or len(endpointSocks5) > 2: |
78 | 96 | parser.error(f"Invalid host option '{values}'") |
79 | 97 |
|
| 98 | + endpoint = HostAction._parseHostPort(parser, endpointSocks5[0]) |
| 99 | + socks5 = None |
| 100 | + if len(endpointSocks5) == 2: |
| 101 | + socks5 = HostAction._parseHostPort(parser, endpointSocks5[1]) |
| 102 | + |
80 | 103 | hosts = getattr(args, self.dest) |
81 | | - hosts.append((vals[0], int(vals[1]))) |
| 104 | + hosts.append(ServerAddress(endpoint, socks5)) |
82 | 105 |
|
83 | 106 |
|
84 | 107 | class UserIdIpAction(Action): |
@@ -106,8 +129,9 @@ def addConnectionAndAuthOptions(parser, forClientServerSetup=False): |
106 | 129 | "-H", |
107 | 130 | "--host", |
108 | 131 | dest="hosts", |
109 | | - help="server name or IP (default: 127.0.0.1:8194). Can be specified multiple times.", |
110 | | - metavar="host:port", |
| 132 | + help="Endpoint host:port and optional SOCKS5 host:port to use to connect" |
| 133 | + " (default: localhost:8194). Can be specified multiple times.", |
| 134 | + metavar="host:port[/socks5Host:socks5Port]", |
111 | 135 | action=HostAction, |
112 | 136 | default=[], |
113 | 137 | ) |
@@ -291,18 +315,37 @@ def createSessionOptions(options): |
291 | 315 | ) |
292 | 316 | else: |
293 | 317 | sessionOptions = blpapi.SessionOptions() |
294 | | - for idx, host in enumerate(options.hosts): |
295 | | - sessionOptions.setServerAddress(host[0], host[1], idx) |
| 318 | + for idx, serverAddress in enumerate(options.hosts): |
| 319 | + socks5Config = ( |
| 320 | + blpapi.Socks5Config( |
| 321 | + serverAddress.socks5.host, int(serverAddress.socks5.port) |
| 322 | + ) |
| 323 | + if serverAddress.socks5 |
| 324 | + else None |
| 325 | + ) |
| 326 | + sessionOptions.setServerAddress( |
| 327 | + serverAddress.endpoint.host, |
| 328 | + int(serverAddress.endpoint.port), |
| 329 | + idx, |
| 330 | + socks5Config, |
| 331 | + ) |
296 | 332 |
|
297 | 333 | if tlsOptions: |
298 | 334 | sessionOptions.setTlsOptions(tlsOptions) |
299 | 335 |
|
300 | 336 | sessionOptions.setSessionIdentityOptions( |
301 | 337 | options.sessionIdentityAuthOptions |
302 | 338 | ) |
| 339 | + |
| 340 | + def serverAddressToString(server: ServerAddress): |
| 341 | + tostr = ":".join(server.endpoint) |
| 342 | + if server.socks5: |
| 343 | + tostr = f"{tostr}/{':'.join(server.socks5)}" |
| 344 | + return tostr |
| 345 | + |
303 | 346 | print( |
304 | 347 | f"Connecting to " |
305 | | - f"{', '.join([h[0] + ':' + str(h[1]) for h in sessionOptions.serverAddresses()])}" |
| 348 | + f"{', '.join([serverAddressToString(h) for h in options.hosts])}" |
306 | 349 | ) |
307 | 350 |
|
308 | 351 | return sessionOptions |
|
0 commit comments