diff --git a/src/meshctrl/session.py b/src/meshctrl/session.py index 1292df7..88c240a 100644 --- a/src/meshctrl/session.py +++ b/src/meshctrl/session.py @@ -46,9 +46,19 @@ class Session(object): ''' def __init__(self, url, user=None, domain=None, password=None, loginkey=None, proxy=None, token=None, ignore_ssl=False, auto_reconnect=False): - if len(url) < 5 or ((not url.startswith('wss://')) and (not url.startswith('ws://'))): + parsed = urllib.parse.urlparse(url) + + if parsed.scheme not in ("wss", "ws"): raise ValueError("Invalid URL") + port = 80 + if parsed.port is None: + if parsed.scheme == "wss": + port = 443 + p = list(parsed) + p[1] = f"{parsed.hostname}:{port}" + url = urllib.parse.urlunparse(p) + if (not url.endswith('/')): url += '/' diff --git a/tests/test_sanity.py b/tests/test_sanity.py index edc265b..9dbc4d9 100644 --- a/tests/test_sanity.py +++ b/tests/test_sanity.py @@ -36,4 +36,21 @@ async def test_ssl(env): except* ssl.SSLCertVerificationError: pass else: - raise Exception("Invalid SSL certificate accepted") \ No newline at end of file + raise Exception("Invalid SSL certificate accepted") + +async def test_urlparse(): + # This tests the url port adding necessitated by python-socks. Our test environment doesn't use 443, so this is just a quick sanity test. + try: + async with meshctrl.Session("wss://localhost", user="unprivileged", password="Not a real password", ignore_ssl=True) as s: + pass + except* TimeoutError: + #We're not running a server, so timeout is our expected outcome + pass + + # This tests our check for wss/ws url schemes + try: + async with meshctrl.Session("https://localhost", user="unprivileged", password="Not a real password", ignore_ssl=True) as s: + pass + except* ValueError: + #We're not running a server, so timeout is our expected outcome + pass \ No newline at end of file