windscribe: more cleaner api to work with cookie

* external storage is supported
* proper login context added
* use existing cookie to perform webui changes

src/config.py

@@ -5,8 +5,6 @@
 import sys
 from pathlib import Path
 
-import httpx
-
 # only run once
 # allows wg-ephemeral to be used as a cronjob
 _ONESHOT: str = os.getenv("ONESHOT", "false")
@@ -34,6 +32,7 @@
 # WS config
 WS_USERNAME: str = os.getenv("WS_USERNAME", "")
 WS_PASSWORD: str = os.getenv("WS_PASSWORD", "")
+WS_COOKIE = Path(os.getenv("WS_COOKIE_PATH", ".")) / "cookie.pkl"
 
 if not all([WS_USERNAME, WS_PASSWORD]):
     print("ENV: WS_USERNAME and WS_PASSWORD need to be set")
@@ -44,7 +43,6 @@
 USERNAME_ID: str = "username"
 PASSWORD_ID: str = "password"
 
-COOKIE_PATH = Path("cookie.pkl")
 
 # fmt: off
 # some qbit config

src/lib/decorators.py

@@ -0,0 +1,21 @@
+import functools
+import logging
+from typing import TYPE_CHECKING, Any, Callable
+
+if TYPE_CHECKING:
+    from ws import Windscribe
+
+
+def login_required(func: Callable):
+    @functools.wraps(func)
+    def inner(*args: Any, **kwargs: Any):
+        obj: "Windscribe" = args[0]
+
+        # if any requests made and not authenticated then login first.
+        if not obj.is_authenticated:
+            logging.warning("rquests is authenticated, sending one.")
+            obj.login()
+
+        return func(*args, **kwargs)
+
+    return inner

src/ws/cookie.py

@@ -14,23 +14,33 @@ def default_cookie() -> Cookies:
 
 
 def load_cookie() -> None | Cookies:
-    """Load existing cookie."""
-    if not config.COOKIE_PATH.exists():
+    """Load existing cookie.
+
+    Read the pickle back and create the cookie object for httpx.
+    """
+    if not config.WS_COOKIE.exists():
         return None
 
     cookie = Cookies()
-    with open(config.COOKIE_PATH, "rb") as ck:
+    with open(config.WS_COOKIE, "rb") as ck:
         cookie_dict = pickle.load(ck)
         for k, v in cookie_dict.items():
             cookie.set(k, v)
-        return cookie
 
+    return cookie
+
+
+def save_cookie(cookie: Cookies) -> None:
+    """Save the cookie to the file for future use.
+
+    Read the cookie data and convert to regular dictinary object so that it can be
+    pickled to a file.
 
-def save_cookie(cookie: Cookies):
-    """Save the cookie to the file for future use"""
+    :param cookie: a cookie object from httpx requests.
+    """
     cookie_dict: dict[str, str] = {}
     for k, v in cookie.items():
         cookie_dict[k] = v
 
-    with open(config.COOKIE_PATH, "wb") as ck:
+    with open(config.WS_COOKIE, "wb") as ck:
         pickle.dump(cookie_dict, ck)

src/ws/ws.py

@@ -14,6 +14,7 @@
 import httpx
 
 import config
+from lib.decorators import login_required
 
 from .cookie import default_cookie, load_cookie, save_cookie
 
@@ -26,33 +27,31 @@ class Csrf(TypedDict):
 
 
 class Windscribe:
-    """
-    windscribe api to enable ephemeral ports.
-    Only works with non 2FA account.
+    """Windscribe api to enable ephemeral ports.
+
+    Only works with non 2FA account (for now).
     """
 
     # pylint: disable=redefined-outer-name
     def __init__(self, username: str, password: str) -> None:
         headers = {
             "origin": config.BASE_URL,
             "referer": config.LOGIN_URL,
-            # pylint: disable=line-too-long
-            # ruff: noqa: E501
-            "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/104.0.0.0 Safari/537.36",
+            "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/104.0.0.0 Safari/537.36",  # ruff: noqa: E501
         }
 
-        self.cookie_found = True
+        self.is_authenticated = True
         cookie = load_cookie()
         if cookie is None:
-            self.cookie_found = False
+            self.is_authenticated = False
             cookie = default_cookie()
 
         self.client = httpx.Client(
             headers=headers, cookies=cookie, timeout=config.REQUEST_TIMEOUT
         )
 
         # we will populate this later in the login call
-        self.csrf: Csrf = self._get_csrf()
+        self.csrf: Csrf = self.get_csrf()
         self.username = username
         self.password = password
 
@@ -66,12 +65,23 @@ def __exit__(self, exc_type, exc_value, traceback) -> None:
         """close httpx session"""
         self.close()
 
-    def _get_csrf(self) -> Csrf:
+    @property
+    def is_authenticated(self) -> bool:
+        """If session is authenticated."""
+        return self._is_authenticated
+
+    @is_authenticated.setter
+    def is_authenticated(self, value: bool) -> None:
+        """Set authentication status."""
+        self._is_authenticated = value
+
+    def get_csrf(self) -> Csrf:
         """windscribe make seperate request to get the csrf token"""
         resp = self.client.post(config.CSRF_URL)
         return resp.json()
 
-    def _renew_csrf(self) -> Csrf:
+    @login_required
+    def renew_csrf(self) -> Csrf:
         """after login windscribe issue new csrf token withing javascript"""
         resp = self.client.get(config.MYACT_URL)
         csrf_time = re.search(r"csrf_time = (?P<ctime>\d+)", resp.text)
@@ -90,8 +100,8 @@ def _renew_csrf(self) -> Csrf:
         self.logger.debug("csrf renewed successfully.")
         return new_csrf
 
-    def _login(self) -> None:
-        """login in to the webpage"""
+    def login(self) -> None:
+        """login in to the webpage."""
         data = {
             "login": 1,
             "upgrade": 0,
@@ -106,9 +116,11 @@ def _login(self) -> None:
         # save the cookie for the future use.
         save_cookie(self.client.cookies)
 
+        self.is_authenticated = True
         self.logger.debug("login successful")
 
-    def _delete_ephm_port(self) -> dict[str, Union[bool, int]]:
+    @login_required
+    def delete_ephm_port(self) -> dict[str, Union[bool, int]]:
         """
         ensure we delete the ephemeral port setting if any available
         """
@@ -122,7 +134,8 @@ def _delete_ephm_port(self) -> dict[str, Union[bool, int]]:
 
         return res
 
-    def _set_matching_port(self) -> int:
+    @login_required
+    def set_matching_port(self) -> int:
         """
         setup matching ephemeral port on WS
         """
@@ -150,15 +163,12 @@ def _set_matching_port(self) -> int:
 
     def setup(self) -> int:
         """perform ephemeral port setup here"""
-        if not self.cookie_found:
-            self._login()
-
         # after login we need to update the csrf token agian,
         # windscribe puts new csrf token in the javascript
-        self.csrf = self._renew_csrf()
+        self.csrf = self.renew_csrf()
 
-        self._delete_ephm_port()
-        return self._set_matching_port()
+        self.delete_ephm_port()
+        return self.set_matching_port()
 
     def close(self) -> None:
         """close httpx session"""