|
23 | 23 | import six
|
24 | 24 | import yaml
|
25 | 25 |
|
26 |
| -from six.moves.urllib.parse import urlencode, quote_plus, urlparse, urlunparse |
| 26 | +from six.moves.urllib.parse import urlencode, urlparse, urlunparse |
27 | 27 | from six import StringIO
|
28 | 28 |
|
29 | 29 | from websocket import WebSocket, ABNF, enableTrace
|
@@ -51,47 +51,13 @@ def __init__(self, configuration, url, headers, capture_all):
|
51 | 51 | like port forwarding can forward different pods' streams to different
|
52 | 52 | channels.
|
53 | 53 | """
|
54 |
| - enableTrace(False) |
55 |
| - header = [] |
56 | 54 | self._connected = False
|
57 | 55 | self._channels = {}
|
58 | 56 | if capture_all:
|
59 | 57 | self._all = StringIO()
|
60 | 58 | else:
|
61 | 59 | self._all = _IgnoredIO()
|
62 |
| - |
63 |
| - # We just need to pass the Authorization, ignore all the other |
64 |
| - # http headers we get from the generated code |
65 |
| - if headers and 'authorization' in headers: |
66 |
| - header.append("authorization: %s" % headers['authorization']) |
67 |
| - |
68 |
| - if headers and 'sec-websocket-protocol' in headers: |
69 |
| - header.append("sec-websocket-protocol: %s" % |
70 |
| - headers['sec-websocket-protocol']) |
71 |
| - else: |
72 |
| - header.append("sec-websocket-protocol: v4.channel.k8s.io") |
73 |
| - |
74 |
| - if url.startswith('wss://') and configuration.verify_ssl: |
75 |
| - ssl_opts = { |
76 |
| - 'cert_reqs': ssl.CERT_REQUIRED, |
77 |
| - 'ca_certs': configuration.ssl_ca_cert or certifi.where(), |
78 |
| - } |
79 |
| - if configuration.assert_hostname is not None: |
80 |
| - ssl_opts['check_hostname'] = configuration.assert_hostname |
81 |
| - else: |
82 |
| - ssl_opts = {'cert_reqs': ssl.CERT_NONE} |
83 |
| - |
84 |
| - if configuration.cert_file: |
85 |
| - ssl_opts['certfile'] = configuration.cert_file |
86 |
| - if configuration.key_file: |
87 |
| - ssl_opts['keyfile'] = configuration.key_file |
88 |
| - |
89 |
| - self.sock = WebSocket(sslopt=ssl_opts, skip_utf8_validation=False) |
90 |
| - if configuration.proxy: |
91 |
| - proxy_url = urlparse(configuration.proxy) |
92 |
| - self.sock.connect(url, header=header, http_proxy_host=proxy_url.hostname, http_proxy_port=proxy_url.port) |
93 |
| - else: |
94 |
| - self.sock.connect(url, header=header) |
| 60 | + self.sock = create_websocket(configuration, url, headers) |
95 | 61 | self._connected = True
|
96 | 62 |
|
97 | 63 | def peek_channel(self, channel, timeout=0):
|
@@ -259,41 +225,86 @@ def close(self, **kwargs):
|
259 | 225 | WSResponse = collections.namedtuple('WSResponse', ['data'])
|
260 | 226 |
|
261 | 227 |
|
262 |
| -def get_websocket_url(url): |
| 228 | +def get_websocket_url(url, query_params=None): |
263 | 229 | parsed_url = urlparse(url)
|
264 | 230 | parts = list(parsed_url)
|
265 | 231 | if parsed_url.scheme == 'http':
|
266 | 232 | parts[0] = 'ws'
|
267 | 233 | elif parsed_url.scheme == 'https':
|
268 | 234 | parts[0] = 'wss'
|
| 235 | + if query_params: |
| 236 | + query = [] |
| 237 | + for key, value in query_params: |
| 238 | + if key == 'command' and isinstance(value, list): |
| 239 | + for command in value: |
| 240 | + query.append((key, command)) |
| 241 | + else: |
| 242 | + query.append((key, value)) |
| 243 | + if query: |
| 244 | + parts[4] = urlencode(query) |
269 | 245 | return urlunparse(parts)
|
270 | 246 |
|
271 | 247 |
|
272 |
| -def websocket_call(configuration, *args, **kwargs): |
| 248 | +def create_websocket(configuration, url, headers=None): |
| 249 | + enableTrace(False) |
| 250 | + |
| 251 | + # We just need to pass the Authorization, ignore all the other |
| 252 | + # http headers we get from the generated code |
| 253 | + header = [] |
| 254 | + if headers and 'authorization' in headers: |
| 255 | + header.append("authorization: %s" % headers['authorization']) |
| 256 | + if headers and 'sec-websocket-protocol' in headers: |
| 257 | + header.append("sec-websocket-protocol: %s" % |
| 258 | + headers['sec-websocket-protocol']) |
| 259 | + else: |
| 260 | + header.append("sec-websocket-protocol: v4.channel.k8s.io") |
| 261 | + |
| 262 | + if url.startswith('wss://') and configuration.verify_ssl: |
| 263 | + ssl_opts = { |
| 264 | + 'cert_reqs': ssl.CERT_REQUIRED, |
| 265 | + 'ca_certs': configuration.ssl_ca_cert or certifi.where(), |
| 266 | + } |
| 267 | + if configuration.assert_hostname is not None: |
| 268 | + ssl_opts['check_hostname'] = configuration.assert_hostname |
| 269 | + else: |
| 270 | + ssl_opts = {'cert_reqs': ssl.CERT_NONE} |
| 271 | + |
| 272 | + if configuration.cert_file: |
| 273 | + ssl_opts['certfile'] = configuration.cert_file |
| 274 | + if configuration.key_file: |
| 275 | + ssl_opts['keyfile'] = configuration.key_file |
| 276 | + |
| 277 | + websocket = WebSocket(sslopt=ssl_opts, skip_utf8_validation=False) |
| 278 | + if configuration.proxy: |
| 279 | + proxy_url = urlparse(configuration.proxy) |
| 280 | + websocket.connect(url, header=header, http_proxy_host=proxy_url.hostname, http_proxy_port=proxy_url.port) |
| 281 | + else: |
| 282 | + websocket.connect(url, header=header) |
| 283 | + return websocket |
| 284 | + |
| 285 | + |
| 286 | +def _configuration(api_client): |
| 287 | + # old generated code's api client has config. new ones has |
| 288 | + # configuration |
| 289 | + try: |
| 290 | + return api_client.configuration |
| 291 | + except AttributeError: |
| 292 | + return api_client.config |
| 293 | + |
| 294 | + |
| 295 | +def websocket_call(api_client, _method, url, **kwargs): |
273 | 296 | """An internal function to be called in api-client when a websocket
|
274 | 297 | connection is required. args and kwargs are the parameters of
|
275 | 298 | apiClient.request method."""
|
276 | 299 |
|
277 |
| - url = args[1] |
| 300 | + url = get_websocket_url(url, kwargs.get("query_params")) |
| 301 | + headers = kwargs.get("headers") |
278 | 302 | _request_timeout = kwargs.get("_request_timeout", 60)
|
279 | 303 | _preload_content = kwargs.get("_preload_content", True)
|
280 | 304 | capture_all = kwargs.get("capture_all", True)
|
281 |
| - headers = kwargs.get("headers") |
282 |
| - |
283 |
| - # Expand command parameter list to indivitual command params |
284 |
| - query_params = [] |
285 |
| - for key, value in kwargs.get("query_params", {}): |
286 |
| - if key == 'command' and isinstance(value, list): |
287 |
| - for command in value: |
288 |
| - query_params.append((key, command)) |
289 |
| - else: |
290 |
| - query_params.append((key, value)) |
291 |
| - |
292 |
| - if query_params: |
293 |
| - url += '?' + urlencode(query_params) |
294 | 305 |
|
295 | 306 | try:
|
296 |
| - client = WSClient(configuration, get_websocket_url(url), headers, capture_all) |
| 307 | + client = WSClient(_configuration(api_client), url, headers, capture_all) |
297 | 308 | if not _preload_content:
|
298 | 309 | return client
|
299 | 310 | client.run_forever(timeout=_request_timeout)
|
|
0 commit comments