Merge pull request #23 from BITNP/swo/get-status

README.md

@@ -29,6 +29,25 @@ pipx install bitsrun
 
 ### CLI
 
+Check login status of your device.
+
+![bitsrun status](https://user-images.githubusercontent.com/32114380/216757172-368d74bc-ad74-4122-9b1f-9568ce0341d3.png)
+
+```text
+Usage: bitsrun status [OPTIONS]
+
+  Check current network login status.
+
+Options:
+  --help  Show this message and exit.
+```
+
+> **Note**: this is the output of `bitsrun status --help`.
+
+Login or logout with your username and password.
+
+![bitsrun login](https://user-images.githubusercontent.com/32114380/216757151-b6e8c620-48b6-4411-ac41-f07b79ef9827.png)
+
 ```text
 Usage: bitsrun login/logout [OPTIONS]
 

bitsrun/cli.py

@@ -1,11 +1,12 @@
 import sys
 from getpass import getpass
-from pprint import pprint
 
 import click
+from rich import print_json
 
 from bitsrun.config import get_config_paths, read_config
-from bitsrun.user import User
+from bitsrun.user import User, get_login_status
+from bitsrun.utils import print_status_table
 
 # A hacky way to specify shared options for multiple click commands:
 # https://stackoverflow.com/questions/40182157/shared-options-and-flags-between-commands
@@ -39,6 +40,25 @@ def config_paths():
     click.echo("\n".join(map(str, get_config_paths())))
 
 
+@cli.command()
+def status():
+    """Check current network login status."""
+    login_status = get_login_status()
+
+    if login_status.get("user_name"):
+        click.echo(
+            click.style("bitsrun: ", fg="green")
+            + f"{login_status['user_name']} ({login_status['online_ip']}) is online"
+        )
+        print_status_table(login_status)
+
+    else:
+        click.echo(
+            click.style("bitsrun: ", fg="cyan")
+            + f"{login_status['online_ip']} is offline"
+        )
+
+
 @cli.command()
 @add_options(_options)
 def login(username, password, verbose):
@@ -58,23 +78,24 @@ def do_action(action, username, password, verbose):
     if username and not password:
         password = getpass(prompt="Please enter your password: ")
 
-    # Try to read username and password from args provided. If none, look for config
-    # files in possible paths. If none, fail and prompt user to provide one.
-    if username and password:
-        user = User(username, password)
-    elif conf := read_config():
-        user = User(**conf[0])
-        if verbose:
-            click.echo(
-                click.style("bitsrun: ", fg="blue")
-                + "Reading config from "
-                + click.style(conf[1], fg="yellow", underline=True)
-            )
-    else:
-        ctx = click.get_current_context()
-        ctx.fail("No username or password provided")
-
     try:
+        # Try to read username and password from args provided. If none, look for config
+        # files in possible paths. If none, fail and prompt user to provide one.
+        if username and password:
+            user = User(username, password)
+        elif conf := read_config():
+            if verbose:
+                click.echo(
+                    click.style("bitsrun: ", fg="blue")
+                    + "Reading config from "
+                    + click.style(conf[1], fg="yellow", underline=True)
+                )
+            user = User(**conf[0])
+        else:
+            ctx = click.get_current_context()
+            ctx.fail("No username or password provided")
+            sys.exit(1)
+
         if action == "login":
             resp = user.login()
             message = f"{user.username} ({resp['online_ip']}) logged in"
@@ -88,7 +109,7 @@ def do_action(action, username, password, verbose):
         # Output direct result of the API response if verbose
         if verbose:
             click.echo(f"{click.style('bitsrun:', fg='cyan')} Response from API:")
-            pprint(resp)
+            print_json(data=resp)
 
         # Handle error from API response. When field `error` is not `ok`, then the
         # login/logout action has likely failed. Hints are provided in the `error_msg`.

bitsrun/models.py

@@ -0,0 +1,32 @@
+from enum import Enum
+from typing import Literal, Optional, TypedDict, Union
+
+
+class Action(Enum):
+    LOGIN = "login"
+    LOGOUT = "logout"
+
+
+class UserResponseType(TypedDict):
+    client_ip: str
+    online_ip: str
+    # Field `error` is also `login_error` when logout action fails
+    error: Union[Literal["login_error"], Literal["ok"]]
+    error_msg: str
+    res: Union[Literal["login_error"], Literal["ok"]]
+    # Field `username` is not present on login fails and all logout scenarios
+    username: Optional[str]
+
+
+class LoginStatusRespType(TypedDict):
+    # Field `error` is `not_online_error` when device is not logged in
+    error: str
+    client_ip: Optional[str]
+    # Field `online_ip` is always present regardless of login status
+    online_ip: str
+    # Below are fields only present when device is logged in
+    sum_bytes: Optional[int]
+    sum_seconds: Optional[int]
+    user_balance: Optional[int]
+    user_name: Optional[str]
+    wallet_balance: Optional[int]

bitsrun/user.py

@@ -1,56 +1,75 @@
 import hmac
 import json
-from enum import Enum
 from hashlib import sha1
-from typing import Dict, Literal, Optional, TypedDict, Union
+from typing import Dict, Optional, Union
 
 import httpx
 
-from bitsrun.utils import fkbase64, parse_homepage, xencode
+from bitsrun.models import Action, LoginStatusRespType, UserResponseType
+from bitsrun.utils import fkbase64, xencode
 
 _API_BASE = "http://10.0.0.55"
 _TYPE_CONST = 1
 _N_CONST = 200
 
 
-class Action(Enum):
-    LOGIN = "login"
-    LOGOUT = "logout"
+def get_login_status(client: Optional[httpx.Client] = None) -> LoginStatusRespType:
+    """Get current login status of the device.
 
+    Note:
+        This function is also used without initializing a `User` instance. As such,
+        the `client` argument is optional and will be initialized if not provided.
 
-class UserResponseType(TypedDict):
-    client_ip: str
-    online_ip: str
-    # Field `error` is also `login_error` when logout action fails
-    error: Union[Literal["login_error"], Literal["ok"]]
-    error_msg: str
-    res: Union[Literal["login_error"], Literal["ok"]]
-    # Field `username` is not present on login fails and all logout scenarios
-    username: Optional[str]
+    Args:
+        client: An optional reused httpx client if provided. Defaults to None.
+
+    Returns:
+        The login status of the current device. If the device is logged in, the
+        `user_name` field will be present. Otherwise, it will be `None`. As such,
+        the presence of `user_name` is used to check if the device is logged in.
+    """
+
+    if not client:
+        client = httpx.Client(base_url=_API_BASE)
+
+    resp = client.get("/cgi-bin/rad_user_info", params={"callback": "jsonp"})
+    return json.loads(resp.text[6:-1])
 
 
 class User:
     def __init__(self, username: str, password: str):
         self.username = username
         self.password = password
 
-        self.ip, self.acid = parse_homepage(api_base=_API_BASE)
+        # Initialize reused httpx client
         self.client = httpx.Client(base_url=_API_BASE)
 
-    def login(self) -> UserResponseType:
-        logged_in_user = self._user_validate()
+        # Get `ac_id` from the redirected login page
+        resp = self.client.get("/", follow_redirects=True)
+        self.acid = resp.url.params.get("ac_id")
 
+        # Check current login status and get device `online_ip`
+        login_status = get_login_status(client=self.client)
+        self.ip = login_status.get("online_ip")
+        self.logged_in_user = login_status.get("user_name")
+
+        # Validate if current logged in user matches the provided username
+        if self.logged_in_user and self.logged_in_user != self.username:
+            raise Exception(
+                f"Current logged in user ({self.logged_in_user}) and "
+                f"yours ({self.username}) does not match"
+            )
+
+    def login(self) -> UserResponseType:
         # Raise exception if device is already logged in
-        if logged_in_user == self.username:
-            raise Exception(f"{logged_in_user}, you are already online")
+        if self.logged_in_user == self.username:
+            raise Exception(f"{self.logged_in_user}, you are already online")
 
         return self._do_action(Action.LOGIN)
 
     def logout(self) -> UserResponseType:
-        logged_in_user = self._user_validate()
-
         # Raise exception if device is not logged in
-        if logged_in_user is None:
+        if self.logged_in_user is None:
             raise Exception("you have already logged out")
 
         return self._do_action(Action.LOGOUT)
@@ -60,43 +79,6 @@ def _do_action(self, action: Action) -> UserResponseType:
         response = self.client.get("/cgi-bin/srun_portal", params=params)
         return json.loads(response.text[6:-1])
 
-    def _get_user_info(self) -> Optional[str]:
-        """Get current logged in user info if exists.
-
-        Returns:
-            The username of the current logged in user if exists.
-        """
-
-        resp = self.client.get("/cgi-bin/rad_user_info")
-        data = resp.text
-
-        if data == "not_online_error":
-            return None
-
-        return data.split(",")[0]
-
-    def _user_validate(self) -> Optional[str]:
-        """Check if current logged in user matches the username provided.
-
-        Raises:
-            Exception: If current logged in user and username provided does not match.
-
-        Returns:
-            The username of the current logged in user if exists.
-        """
-
-        logged_in_user = self._get_user_info()
-
-        # Raise exception only if username exists on this IP and
-        # command line arguments provided another username
-        if logged_in_user and logged_in_user != self.username:
-            raise Exception(
-                f"Current logged in user ({logged_in_user}) and "
-                f"yours ({self.username}) does not match"
-            )
-
-        return logged_in_user
-
     def _get_token(self) -> str:
         params = {"callback": "jsonp", "username": self.username, "ip": self.ip}
         response = self.client.get("/cgi-bin/get_challenge", params=params)

bitsrun/utils.py

@@ -1,53 +1,45 @@
 import math
 from base64 import b64encode
-from html.parser import HTMLParser
-from typing import Tuple
 
-import httpx
+from humanize import naturaldelta, naturalsize
+from rich import box
+from rich.console import Console
+from rich.table import Table
 
+from bitsrun.models import LoginStatusRespType
 
-def parse_homepage(api_base: str) -> Tuple[str, str]:
-    """Parse homepage of 10.0.0.55 and get the acid + ip of current session.
 
-    Raises:
-        Exception: Throw exception if acid not present in the redirected URL.
-        Exception: Throw exception if response text does not contain IP.
+def print_status_table(login_status: LoginStatusRespType) -> None:
+    """Print the login status table to the console if logged in.
 
-    Returns:
-        A tuple of (ip, acid) of the current session.
-    """
-
-    res = httpx.get(api_base, follow_redirects=True)
-
-    # ac_id appears in the url query parameter of the redirected URL
-    ac_id = res.url.params.get("ac_id")
+    You should get something like this:
 
-    if not ac_id:
-        raise Exception("failed to get acid")
-
-    # ip appears in the response HTML
-    class IPParser(HTMLParser):
-        def __init__(self, *args, **kwargs):
-            super().__init__(*args, **kwargs)
-            self.ip = None
+    ┌──────────────┬──────────────┬──────────────┬──────────────┐
+    │ Traffic Used │ Online Time  │ User Balance │ Wallet       │
+    ├──────────────┼──────────────┼──────────────┼──────────────┤
+    │ 879.3 MiB    │ 3 hours      │ 10.00        │ 0.00         │
+    └──────────────┴──────────────┴──────────────┴──────────────┘
+    """
 
-        def handle_starttag(self, tag, attrs):
-            if tag == "input":
-                attr_dict = dict(attrs)
-                if attr_dict.get("name") == "user_ip":
-                    self.ip = attr_dict["value"]
+    if not login_status.get("user_name"):
+        return
 
-        def feed(self, *args, **kwargs):
-            super().feed(*args, **kwargs)
-            return self.ip
+    table = Table(box=box.SQUARE)
 
-    parser = IPParser()
-    ip = parser.feed(res.text)
+    table.add_column("Traffic Used", style="magenta", width=12)
+    table.add_column("Online Time", style="yellow", width=12)
+    table.add_column("User Balance", style="green", width=12)
+    table.add_column("Wallet", style="blue", width=12)
 
-    if not ip:
-        raise Exception("failed to get ip")
+    table.add_row(
+        naturalsize(login_status.get("sum_bytes", 0), binary=True),  # type: ignore
+        naturaldelta(login_status.get("sum_seconds", 0)),  # type: ignore
+        f"{login_status.get('user_balance', 0):0.2f}",
+        f"{login_status.get('wallet_balance', 0):0.2f}",
+    )
 
-    return ip, ac_id[0]
+    console = Console()
+    console.print(table)
 
 
 def fkbase64(raw_s: str) -> str: