From 9bdf163a6859fffd2bca9e141aee58bd7ed2ae6d Mon Sep 17 00:00:00 2001 From: spencerwooo Date: Tue, 24 Jan 2023 22:09:21 +0800 Subject: [PATCH 1/6] feat: better error handling and type defs --- bitsrun/cli.py | 56 +++++++++++++++++++++++++++-------------------- bitsrun/config.py | 13 +++++++---- bitsrun/user.py | 39 ++++++++++++++++++++------------- 3 files changed, 65 insertions(+), 43 deletions(-) diff --git a/bitsrun/cli.py b/bitsrun/cli.py index f93c2bf..41d8a8d 100644 --- a/bitsrun/cli.py +++ b/bitsrun/cli.py @@ -9,10 +9,9 @@ # A hacky way to specify shared options for multiple click commands: # https://stackoverflow.com/questions/40182157/shared-options-and-flags-between-commands _options = [ - click.option("-u", "--username", help="Username.", required=False), - click.option("-p", "--password", help="Password.", required=False), + click.option("-u", "--username", help="Your username.", required=False), + click.option("-p", "--password", help="Your password.", required=False), click.option("-v", "--verbose", is_flag=True, help="Verbosely echo API response."), - click.option("-s", "--silent", is_flag=True, help="Silent, no output to stdout."), ] @@ -41,52 +40,61 @@ def config_paths(): @cli.command() @add_options(_options) -def login(username, password, verbose, silent): +def login(username, password, verbose): """Log into the BIT network.""" - do_action("login", username, password, verbose, silent) + do_action("login", username, password, verbose) @cli.command() @add_options(_options) -def logout(username, password, verbose, silent): +def logout(username, password, verbose): """Log out of the BIT network.""" - do_action("logout", username, password, verbose, silent) + do_action("logout", username, password, verbose) -def do_action(action, username, password, verbose, silent): - """Log in/out the BIT network.""" +def do_action(action, username, password, verbose): + # Support reading password from stdin when not passed via `--password` 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) + 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/password provided.") + ctx.fail("No username/password provided") try: if action == "login": res = user.login() - - # Output login result by default if not silent - if not silent: - click.echo(f"{res.get('username')} ({res.get('online_ip')}) logged in") - elif action == "logout": res = user.logout() - - # Output logout result by default if not silent - if not silent: - click.echo(f"{res.get('online_ip')} logged out") - else: # Should not reach here, but just in case - raise ValueError(f"unknown action `{action}`") + raise ValueError(f"Unknown action `{action}`") - # Output direct result of response if verbose + # Output direct result of the API response if verbose if verbose: - click.echo(f"{click.style('info:', fg='blue')} {res}") + click.echo(f"{click.style('bitsrun:', fg='blue')} {res}") + + # Handle error from API response. When field `error` is not `ok`, then the + # login/logout action has likely failed. + if res["error"] != "ok": + raise Exception(res["error"]) + + click.echo( + click.style("bitsrun: ", fg="green") + + f"{res['username']} ({res['online_ip']}) logged in" + ) except Exception as e: click.echo(f"{click.style('error:', fg='red')} {e}") diff --git a/bitsrun/config.py b/bitsrun/config.py index d6865d1..af73ed2 100644 --- a/bitsrun/config.py +++ b/bitsrun/config.py @@ -2,7 +2,7 @@ from os import getenv from pathlib import Path from sys import platform -from typing import Optional, Tuple +from typing import Optional, Tuple, TypedDict from platformdirs import site_config_path, user_config_path @@ -49,7 +49,12 @@ def get_config_paths() -> map: return map(lambda path: path / "bit-user.json", paths) -def read_config() -> Optional[Tuple[str, str]]: +class ConfigType(TypedDict): + username: str + password: str + + +def read_config() -> Optional[Tuple[ConfigType, str]]: """Read config from the first available config file with name `bit-user.json`. The config file should be a JSON file with the following structure: @@ -59,7 +64,7 @@ def read_config() -> Optional[Tuple[str, str]]: ``` Returns: - A tuple of (username, password) if the config file is found. + A tuple of (config, path to config file) if the config file is found. """ paths = get_config_paths() @@ -67,7 +72,7 @@ def read_config() -> Optional[Tuple[str, str]]: try: with open(path) as f: data = json.loads(f.read()) - return data["username"], data["password"] + return data, str(path) except Exception: continue return None diff --git a/bitsrun/user.py b/bitsrun/user.py index 0c2fc9d..c3d86c0 100644 --- a/bitsrun/user.py +++ b/bitsrun/user.py @@ -2,15 +2,15 @@ import json from enum import Enum from hashlib import sha1 -from typing import Dict, Optional, Union +from typing import Dict, Literal, Optional, TypedDict, Union from requests import Session from bitsrun.utils import fkbase64, parse_homepage, xencode -API_BASE = "http://10.0.0.55" -TYPE_CONST = 1 -N_CONST = 200 +_API_BASE = "http://10.0.0.55" +_TYPE_CONST = 1 +_N_CONST = 200 class Action(Enum): @@ -18,15 +18,24 @@ class Action(Enum): LOGOUT = "logout" +class UserResponseType(TypedDict): + client_ip: str + online_ip: str + error: Union[Literal["login_error"], Literal["ok"]] + error_msg: str + res: Union[Literal["login_error"], Literal["ok"]] + username: Optional[str] + + 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) + self.ip, self.acid = parse_homepage(api_base=_API_BASE) self.session = Session() - def login(self) -> Dict[str, Union[str, int]]: + def login(self) -> UserResponseType: logged_in_user = self._user_validate() # Raise exception if device is already logged in @@ -35,7 +44,7 @@ def login(self) -> Dict[str, Union[str, int]]: return self._do_action(Action.LOGIN) - def logout(self) -> Dict[str, Union[str, int]]: + def logout(self) -> UserResponseType: logged_in_user = self._user_validate() # Raise exception if device is not logged in @@ -44,9 +53,9 @@ def logout(self) -> Dict[str, Union[str, int]]: return self._do_action(Action.LOGOUT) - def _do_action(self, action: Action) -> Dict[str, Union[str, int]]: + def _do_action(self, action: Action) -> UserResponseType: params = self._make_params(action) - response = self.session.get(API_BASE + "/cgi-bin/srun_portal", params=params) + response = self.session.get(_API_BASE + "/cgi-bin/srun_portal", params=params) return json.loads(response.text[6:-1]) def _get_user_info(self) -> Optional[str]: @@ -56,7 +65,7 @@ def _get_user_info(self) -> Optional[str]: The username of the current logged in user if exists. """ - resp = self.session.get(API_BASE + "/cgi-bin/rad_user_info") + resp = self.session.get(_API_BASE + "/cgi-bin/rad_user_info") data = resp.text if data == "not_online_error": @@ -88,7 +97,7 @@ def _user_validate(self) -> Optional[str]: def _get_token(self) -> str: params = {"callback": "jsonp", "username": self.username, "ip": self.ip} - response = self.session.get(API_BASE + "/cgi-bin/get_challenge", params=params) + response = self.session.get(_API_BASE + "/cgi-bin/get_challenge", params=params) result = json.loads(response.text[6:-1]) return result["challenge"] @@ -101,8 +110,8 @@ def _make_params(self, action: Action) -> Dict[str, Union[int, str]]: "action": action.value, "ac_id": self.acid, "ip": self.ip, - "type": TYPE_CONST, - "n": N_CONST, + "type": _TYPE_CONST, + "n": _N_CONST, } data = { @@ -123,8 +132,8 @@ def _make_params(self, action: Action) -> Dict[str, Union[int, str]]: hmd5, self.acid, self.ip, - N_CONST, - TYPE_CONST, + _N_CONST, + _TYPE_CONST, info, ).encode() ).hexdigest() From 2ba9fd25ab6dfd7d7e5f02fc281fcab0948db054 Mon Sep 17 00:00:00 2001 From: spencerwooo Date: Tue, 24 Jan 2023 22:11:56 +0800 Subject: [PATCH 2/6] docs: update readme and bump version --- README.md | 5 ++--- pyproject.toml | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index fd18194..be89fdf 100644 --- a/README.md +++ b/README.md @@ -29,10 +29,9 @@ Usage: bitsrun login/logout [OPTIONS] Log into or out of the BIT network. Options: - -u, --username TEXT Username. - -p, --password TEXT Password. + -u, --username TEXT Your username. + -p, --password TEXT Your password. -v, --verbose Verbosely echo API response. - -s, --silent Silent, no output to stdout. --help Show this message and exit. ``` diff --git a/pyproject.toml b/pyproject.toml index efa985e..d60e2ce 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "bitsrun" -version = "3.2.4" +version = "3.3.0" description = "A headless login / logout script for 10.0.0.55" authors = ["spencerwooo "] license = "MIT" From a4f00345d3bc563ca1885e0dcf4100df1a37be85 Mon Sep 17 00:00:00 2001 From: spencerwooo Date: Tue, 24 Jan 2023 22:25:23 +0800 Subject: [PATCH 3/6] fix: fallback to user provided username if None --- bitsrun/cli.py | 9 ++++++--- bitsrun/user.py | 2 ++ 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/bitsrun/cli.py b/bitsrun/cli.py index 41d8a8d..4813d7d 100644 --- a/bitsrun/cli.py +++ b/bitsrun/cli.py @@ -1,5 +1,6 @@ import sys from getpass import getpass +from pprint import pprint import click @@ -84,7 +85,9 @@ 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='blue')} {res}") + click.echo(f"{click.style('bitsrun:', fg='cyan')} Response from API:") + # click.echo(res) + pprint(res) # Handle error from API response. When field `error` is not `ok`, then the # login/logout action has likely failed. @@ -93,11 +96,11 @@ def do_action(action, username, password, verbose): click.echo( click.style("bitsrun: ", fg="green") - + f"{res['username']} ({res['online_ip']}) logged in" + + f"{res.get('username', user.username)} ({res['online_ip']}) logged in" ) except Exception as e: - click.echo(f"{click.style('error:', fg='red')} {e}") + click.echo(f"{click.style('error:', fg='red')} {e}", err=True) # Throw with error code 1 for scripts to pick up error state sys.exit(1) diff --git a/bitsrun/user.py b/bitsrun/user.py index c3d86c0..80a59fc 100644 --- a/bitsrun/user.py +++ b/bitsrun/user.py @@ -21,9 +21,11 @@ class Action(Enum): 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"]] + # On login fails and all logout scenarios, field `username` is not present username: Optional[str] From 5761a863c5f6a0ceeaac5fb8cafbe80d24d57738 Mon Sep 17 00:00:00 2001 From: spencerwooo Date: Tue, 24 Jan 2023 22:43:21 +0800 Subject: [PATCH 4/6] fix: logout on success prints login --- bitsrun/cli.py | 28 ++++++++++++++++------------ bitsrun/user.py | 2 +- 2 files changed, 17 insertions(+), 13 deletions(-) diff --git a/bitsrun/cli.py b/bitsrun/cli.py index 4813d7d..c3c15d8 100644 --- a/bitsrun/cli.py +++ b/bitsrun/cli.py @@ -72,13 +72,15 @@ def do_action(action, username, password, verbose): ) else: ctx = click.get_current_context() - ctx.fail("No username/password provided") + ctx.fail("No username or password provided") try: if action == "login": - res = user.login() + resp = user.login() + message = f"{resp['username']} ({resp['online_ip']}) logged in" elif action == "logout": - res = user.logout() + resp = user.logout() + message = f"{resp['online_ip']} logged out" else: # Should not reach here, but just in case raise ValueError(f"Unknown action `{action}`") @@ -86,20 +88,22 @@ 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:") - # click.echo(res) - pprint(res) + pprint(resp, indent=4) # Handle error from API response. When field `error` is not `ok`, then the - # login/logout action has likely failed. - if res["error"] != "ok": - raise Exception(res["error"]) + # login/logout action has likely failed. Hints are provided in the `error_msg`. + if resp["error"] != "ok": + raise Exception( + resp["error_msg"] + if resp["error_msg"] + else "Action failed, use --verbose for more info" + ) - click.echo( - click.style("bitsrun: ", fg="green") - + f"{res.get('username', user.username)} ({res['online_ip']}) logged in" - ) + # Print success message + click.echo(f"{click.style('bitsrun:', fg='green')} {message}") except Exception as e: + # Exception is caught and printed to stderr click.echo(f"{click.style('error:', fg='red')} {e}", err=True) # Throw with error code 1 for scripts to pick up error state sys.exit(1) diff --git a/bitsrun/user.py b/bitsrun/user.py index 80a59fc..e5c4faf 100644 --- a/bitsrun/user.py +++ b/bitsrun/user.py @@ -25,7 +25,7 @@ class UserResponseType(TypedDict): error: Union[Literal["login_error"], Literal["ok"]] error_msg: str res: Union[Literal["login_error"], Literal["ok"]] - # On login fails and all logout scenarios, field `username` is not present + # Field `username` is not present on login fails and all logout scenarios username: Optional[str] From 742d7dd098f12589811259c6b760f273f3a3a7ca Mon Sep 17 00:00:00 2001 From: spencerwooo Date: Tue, 24 Jan 2023 22:46:12 +0800 Subject: [PATCH 5/6] fix: pprint output as default style --- bitsrun/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bitsrun/cli.py b/bitsrun/cli.py index c3c15d8..fd0835d 100644 --- a/bitsrun/cli.py +++ b/bitsrun/cli.py @@ -88,7 +88,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, indent=4) + pprint(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`. From 321c372122294cc38305e118cfd2789d50d43702 Mon Sep 17 00:00:00 2001 From: spencerwooo Date: Tue, 24 Jan 2023 22:48:47 +0800 Subject: [PATCH 6/6] fix: throw better exception on username undefined --- bitsrun/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bitsrun/cli.py b/bitsrun/cli.py index fd0835d..50362ff 100644 --- a/bitsrun/cli.py +++ b/bitsrun/cli.py @@ -77,7 +77,7 @@ def do_action(action, username, password, verbose): try: if action == "login": resp = user.login() - message = f"{resp['username']} ({resp['online_ip']}) logged in" + message = f"{user.username} ({resp['online_ip']}) logged in" elif action == "logout": resp = user.logout() message = f"{resp['online_ip']} logged out"