This commit is contained in:
coletdjnz 2024-07-14 12:24:04 +12:00
parent a0ca9dc2e4
commit 55540caf6b
No known key found for this signature in database
GPG Key ID: 91984263BB39894A

View File

@ -7,6 +7,7 @@
import random import random
import ssl import ssl
import threading import threading
import time
from http.server import BaseHTTPRequestHandler from http.server import BaseHTTPRequestHandler
from socketserver import BaseRequestHandler, ThreadingTCPServer from socketserver import BaseRequestHandler, ThreadingTCPServer
@ -160,11 +161,15 @@ def handle(self):
self.request.settimeout(None) self.request.settimeout(None)
protocol = websockets.ServerProtocol() protocol = websockets.ServerProtocol()
connection = websockets.sync.server.ServerConnection(socket=self.request, protocol=protocol, close_timeout=10) connection = websockets.sync.server.ServerConnection(socket=self.request, protocol=protocol, close_timeout=10)
try:
connection.handshake() connection.handshake()
for message in connection: for message in connection:
if message == 'proxy_info': if message == 'proxy_info':
connection.send(json.dumps(self.proxy_info)) connection.send(json.dumps(self.proxy_info))
connection.close() except Exception as e:
print(f'Error in websocket proxy: {e}')
finally:
connection.close(code=1001)
class WebSocketSecureProxyHandler(WebSocketProxyHandler): class WebSocketSecureProxyHandler(WebSocketProxyHandler):
@ -244,13 +249,14 @@ def proxy_server(proxy_server_class, request_handler, bind_ip=None, **proxy_serv
server_thread.join() server_thread.join()
class HTTPProxyTestContext(abc.ABC): class HTTPProxyTestContext(abc.ABC):
REQUEST_HANDLER_CLASS = None REQUEST_HANDLER_CLASS = None
REQUEST_PROTO = None REQUEST_PROTO = None
def http_server(self, server_class, *args, **kwargs): def http_server(self, server_class, *args, **kwargs):
return proxy_server(server_class, self.REQUEST_HANDLER_CLASS, *args, **kwargs) server = proxy_server(server_class, self.REQUEST_HANDLER_CLASS, *args, **kwargs)
time.sleep(1) # ensure server is up
return server
@abc.abstractmethod @abc.abstractmethod
def proxy_info_request(self, handler, target_domain=None, target_port=None, **req_kwargs) -> dict: def proxy_info_request(self, handler, target_domain=None, target_port=None, **req_kwargs) -> dict: