daanselen/pylibmeshctrl
1
0
Fork 0
forked from Narcissus/pylibmeshctrl
Code Issues Pull Requests Activity

Fixed auto-reconnect for proxy and created tests for auto-reconnect

Browse Source
This commit is contained in:
Josiah Baldwin
2024-12-10 12:49:50 -08:00
parent 0a657cee48
commit ab2a4c40bc
5 changed files with 77 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
src/meshctrl/session.py
View File

@@ -140,13 +140,9 @@ class Session(object):
token = self._token if self._token else b""
headers['x-meshauth'] = (base64.b64encode(self._user.encode()) + b',' + base64.b64encode(self._password.encode()) + token).decode()
if self._proxy:
proxy = Proxy.from_url(self._proxy)
parsed = urllib.parse.urlparse(self.url)
options["sock"] = await proxy.connect(dest_host=parsed.hostname, dest_port=parsed.port)
options["additional_headers"] = headers
async for websocket in websockets.asyncio.client.connect(self.url, process_exception=util._process_websocket_exception, **options):
async for websocket in util.proxy_connect(self.url, proxy_url=self._proxy, process_exception=util._process_websocket_exception, **options):
self.alive = True
self._socket_open.set()
try:
@@ -154,10 +150,10 @@ class Session(object):
tg.create_task(self._listen_data_task(websocket))
tg.create_task(self._send_data_task(websocket))
except* websockets.ConnectionClosed as e:
self._socket_open.clear()
if not self.auto_reconnect:
raise
except* Exception as eg:
self.alive = False
self._socket_open.clear()

7
src/meshctrl/tunnel.py
View File

@@ -70,12 +70,7 @@ class Tunnel(object):
self.url = self._session.url.replace('/control.ashx', '/meshrelay.ashx?browser=1&p=' + str(self._protocol) + '&nodeid=' + self.node_id + '&id=' + self._tunnel_id + '&auth=' + authcookie["cookie"])
if self._session._proxy:
proxy = Proxy.from_url(self._session._proxy)
parsed = urllib.parse.urlparse(self.url)
options["sock"] = await proxy.connect(dest_host=parsed.hostname, dest_port=parsed.port)
async for websocket in websockets.asyncio.client.connect(self.url, process_exception=util._process_websocket_exception, **options):
async for websocket in util.proxy_connect(self.url, proxy_url=self._session._proxy, process_exception=util._process_websocket_exception, **options):
self.alive = True
self._socket_open.set()
try:

37
src/meshctrl/util.py
View File

@@ -9,6 +9,9 @@ import re
import websockets
import ssl
import functools
import urllib
import python_socks
from python_socks.async_.asyncio import Proxy
from . import exceptions
def _encode_cookie(o, key):
@@ -139,20 +142,36 @@ def compare_dict(dict1, dict2):
def _check_socket(f):
@functools.wraps(f)
async def wrapper(self, *args, **kwargs):
await asyncio.wait_for(self.initialized.wait(), 10)
if not self.alive and self._main_loop_error is not None:
raise self._main_loop_error
elif not self.alive:
raise exceptions.SocketError("Socket Closed")
return await f(self, *args, **kwargs)
try:
async with asyncio.TaskGroup() as tg:
tg.create_task(asyncio.wait_for(self.initialized.wait(), 10))
tg.create_task(asyncio.wait_for(self._socket_open.wait(), 10))
finally:
if not self.alive and self._main_loop_error is not None:
raise self._main_loop_error
elif not self.alive:
raise exceptions.SocketError("Socket Closed")
return await f(self, *args, **kwargs)
return wrapper
def _process_websocket_exception(exc):
tmp = websockets.asyncio.client.process_exception(exc)
# SSLVerification error is a subclass of OSError, but doesn't make sense no retry, so we need to handle it separately.
# SSLVerification error is a subclass of OSError, but doesn't make sense to retry, so we need to handle it separately.
if isinstance(exc, (ssl.SSLCertVerificationError, TimeoutError)):
return exc
if isinstance(exc, python_socks._errors.ProxyError):
return None
return tmp
class Proxy(object):
pass
class proxy_connect(websockets.asyncio.client.connect):
def __init__(self,*args, proxy_url=None, **kwargs):
self.proxy = None
if proxy_url is not None:
self.proxy = Proxy.from_url(proxy_url)
super().__init__(*args, **kwargs)
async def create_connection(self, *args, **kwargs):
if self.proxy is not None:
parsed = urllib.parse.urlparse(self.uri)
self.connection_kwargs["sock"] = await self.proxy.connect(dest_host=parsed.hostname, dest_port=parsed.port)
return await super().create_connection(*args, **kwargs)