diff options
| -rw-r--r-- | openbb_terminal/account/account_controller.py | 109 | ||||
| -rw-r--r-- | openbb_terminal/miscellaneous/i18n/en.yml | 4 | ||||
| -rw-r--r-- | openbb_terminal/session/hub_model.py | 136 | ||||
| -rw-r--r-- | tests/openbb_terminal/account/test_account_controller.py | 92 | ||||
| -rw-r--r-- | tests/openbb_terminal/account/txt/test_account_controller/test_print_help.txt | 5 | ||||
| -rw-r--r-- | tests/openbb_terminal/session/test_hub_model.py | 136 |
6 files changed, 479 insertions, 3 deletions
diff --git a/openbb_terminal/account/account_controller.py b/openbb_terminal/account/account_controller.py index eb27adf79c5..88b3208f97b 100644 --- a/openbb_terminal/account/account_controller.py +++ b/openbb_terminal/account/account_controller.py @@ -7,7 +7,10 @@ from typing import Dict, List, Optional from prompt_toolkit.completion import NestedCompleter -from openbb_terminal import feature_flags as obbff +from openbb_terminal import ( + feature_flags as obbff, + keys_model, +) from openbb_terminal.account.account_model import get_diff, get_routines_info from openbb_terminal.account.account_view import display_routines_list from openbb_terminal.core.config.paths import USER_ROUTINES_DIRECTORY @@ -39,6 +42,9 @@ class AccountController(BaseController): "upload", "download", "delete", + "generate", + "show", + "revoke", ] PATH = "/account/" @@ -91,6 +97,11 @@ class AccountController(BaseController): mt.add_cmd("download") mt.add_cmd("delete") mt.add_raw("\n") + mt.add_info("_personal_access_token_") + mt.add_cmd("generate") + mt.add_cmd("show") + mt.add_cmd("revoke") + mt.add_raw("\n") mt.add_info("_authentication_") mt.add_cmd("logout") console.print(text=mt.menu_text, menu="Account") @@ -426,3 +437,99 @@ class AccountController(BaseController): self.update_runtime_choices() else: console.print("[info]Aborted.[/info]") + + @log_start_end(log=logger) + def call_generate(self, other_args: List[str]) -> None: + """Process generate command.""" + parser = argparse.ArgumentParser( + add_help=False, + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + prog="generate", + description="Generate an OpenBB Personal Access Token.", + ) + parser.add_argument( + "-d", + "--days", + dest="days", + help="Number of days the token will be valid", + type=check_positive, + default=30, + ) + parser.add_argument( + "-s", + "--save", + dest="save", + default=False, + help="Save the token to the keys", + action="store_true", + ) + ns_parser = self.parse_known_args_and_warn(parser, other_args) + if ns_parser: + i = console.input( + "[bold yellow]This will revoke any token that was previously generated." + "\nThis action is irreversible.[/bold yellow]" + "\nAre you sure you want to generate a new token? (y/n): " + ) + if i.lower() not in ["y", "yes"]: + console.print("\n[info]Aborted.[/info]") + return + + response = Hub.generate_personal_access_token( + auth_header=User.get_auth_header(), days=ns_parser.days + ) + if response and response.status_code == 200: + token = response.json().get("token", "") + if token: + console.print(f"\n[info]Token:[/info] {token}\n") + + save_to_keys = False + if not ns_parser.save: + save_to_keys = console.input( + "Would you like to save the token to the keys? (y/n): " + ).lower() in ["y", "yes"] + + if save_to_keys or ns_parser.save: + keys_model.set_openbb_personal_access_token( + key=token, persist=True, show_output=True + ) + + @log_start_end(log=logger) + def call_show(self, other_args: List[str]) -> None: + """Process show command.""" + parser = argparse.ArgumentParser( + add_help=False, + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + prog="show", + description="Show your current OpenBB Personal Access Token.", + ) + ns_parser = self.parse_known_args_and_warn(parser, other_args) + if ns_parser: + response = Hub.get_personal_access_token(auth_header=User.get_auth_header()) + if response and response.status_code == 200: + token = response.json().get("token", "") + if token: + console.print(f"[info]Token:[/info] {token}") + + @log_start_end(log=logger) + def call_revoke(self, other_args: List[str]) -> None: + """Process revoke command.""" + parser = argparse.ArgumentParser( + add_help=False, + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + prog="revoke", + description="Revoke your current OpenBB Personal Access Token.", + ) + ns_parser = self.parse_known_args_and_warn(parser, other_args) + if ns_parser: + i = console.input( + "[bold red]This action is irreversible![/bold red]\n" + "Are you sure you want to revoke your token? (y/n): " + ) + if i.lower() in ["y", "yes"]: + response = Hub.revoke_personal_access_token( + auth_header=User.get_auth_header() + ) + if response and response.status_code in [200, 202]: + console.print("[info]Token revoked.[/info]") + else: + console.print("[info]Aborted.[/info]") diff --git a/openbb_terminal/miscellaneous/i18n/en.yml b/openbb_terminal/miscellaneous/i18n/en.yml index 412ec954128..7a0e637fef5 100644 --- a/openbb_terminal/miscellaneous/i18n/en.yml +++ b/openbb_terminal/miscellaneous/i18n/en.yml @@ -20,6 +20,10 @@ en: _main_menu_: Main menu account/_authentication_: Authentication account/logout: Log out from OpenBB account + account/_personal_access_token_: OpenBB Personal Access Token + account/generate: Generate a Personal Access Token + account/show: Shows the current Personal Access Token + account/revoke: Revoke the Personal Access Token account/_info_: Cloud storage of keys, settings and feature flags account/sync: Turns the cloud synchronization on/off account/pull: Pull data from cloud diff --git a/openbb_terminal/session/hub_model.py b/openbb_terminal/session/hub_model.py index d2b29a45aaf..eff3e78d12d 100644 --- a/openbb_terminal/session/hub_model.py +++ b/openbb_terminal/session/hub_model.py @@ -1,3 +1,4 @@ +import json from typing import Dict, Optional import requests @@ -539,3 +540,138 @@ def list_routines( except Exception: console.print("[red]Failed to list your routines.[/red]") return None + + +def generate_personal_access_token( + auth_header: str, base_url: str = BASE_URL, timeout: int = TIMEOUT, days: int = 30 +) -> Optional[requests.Response]: + """ + Generate an OpenBB Personal Access Token. + + Parameters + ---------- + auth_header : str + The authorization header, e.g. "Bearer <token>". + base_url : str + The base url, by default BASE_URL + timeout : int + The timeout, by default TIMEOUT + days : int + The number of days the token should be valid for. + + Returns + ------- + Optional[requests.Response] + """ + + url = f"{base_url}/sdk/token" + + payload = json.dumps({"days": days}) + headers = { + "Authorization": auth_header, + "Content-Type": "application/json", + } + + try: + response = requests.put(url=url, headers=headers, data=payload, timeout=timeout) + + if response.status_code != 200: + console.print("[red]Failed to generate personal access token.[/red]") + + return response + + except requests.exceptions.ConnectionError: + console.print(f"\n{CONNECTION_ERROR_MSG}") + return None + except requests.exceptions.Timeout: + console.print(f"\n{CONNECTION_TIMEOUT_MSG}") + return None + except Exception: + console.print("[red]Failed to generate personal access token.[/red]") + return None + + +def get_personal_access_token( + auth_header: str, base_url: str = BASE_URL, timeout: int = TIMEOUT +) -> Optional[requests.Response]: + """ + Show the user's OpenBB Personal Access Token. + + Parameters + ---------- + auth_header : str + The authorization header, e.g. "Bearer <token>". + base_url : str + The base url, by default BASE_URL + timeout : int + The timeout, by default TIMEOUT + + Returns + ------- + Optional[requests.Response] + """ + + url = f"{base_url}/sdk/token" + + headers = {"Authorization": auth_header} + + try: + response = requests.get(url=url, headers=headers, timeout=timeout) + + if response.status_code != 200: + console.print("[red]Failed to get personal access token.[/red]") + + return response + + except requests.exceptions.ConnectionError: + console.print(f"\n{CONNECTION_ERROR_MSG}") + return None + except requests.exceptions.Timeout: + console.print(f"\n{CONNECTION_TIMEOUT_MSG}") + return None + except Exception: + console.print("[red]Failed to get personal access token.[/red]") + return None + + +def revoke_personal_access_token( + auth_header: str, base_url: str = BASE_URL, timeout: int = TIMEOUT +) -> Optional[requests.Response]: + """ + Delete the user's OpenBB Personal Access Token. + + Parameters + ---------- + auth_header : str + The authorization header, e.g. "Bearer <token>". + base_url : str + The base url, by default BASE_URL + timeout : int + The timeout, by default TIMEOUT + + Returns + ------- + Optional[requests.Response] + """ + + url = f"{base_url}/sdk/token" + + headers = {"Authorization": auth_header} + + try: + response = requests.delete(url=url, headers=headers, timeout=timeout) + + if response.status_code not in [200, 202]: + console.print("[red]Failed to revoke personal access token.[/red]") + + return response + + except requests.exceptions.ConnectionError: + console.print(f"\n{CONNECTION_ERROR_MSG}") + return None + except requests.exceptions.Timeout: + console.print(f"\n{CONNECTION_TIMEOUT_MSG}") + return None + except Exception: + console.print("[red]Failed to revoke personal access token.[/red]") + return None diff --git a/tests/openbb_terminal/account/test_account_controller.py b/tests/openbb_terminal/account/test_account_controller.py index 71e883a34ba..7f183301779 100644 --- a/tests/openbb_terminal/account/test_account_controller.py +++ b/tests/openbb_terminal/account/test_account_controller.py @@ -28,7 +28,7 @@ CONFIGS = { "USER_DATA_DIRECTORY": "some/path/to/user/data", }, "features_keys": { - "API_KEY_ALPHAVANTAGE": "test_av", + "API_KEY_ALPHAVANTAGE": "test_av", # pragma: allowlist secret "API_FRED_KEY": "test_fred", }, } @@ -258,6 +258,9 @@ def test_call_func_expect_queue(expected_queue, func, queue): "call_upload", "call_download", "call_delete", + "call_generate", + "call_show", + "call_revoke", ], ) def test_call_func_no_parser(func, mocker): @@ -538,7 +541,7 @@ def test_call_download(mocker): @pytest.mark.skip( reason="We should add a `-y or -f` option to make that easier to test" ) -def test_call_delete(mocker): +def test_call_delete(mocker, monkeypatch): controller = account_controller.AccountController(queue=None) path_controller = "openbb_terminal.account.account_controller" @@ -549,6 +552,9 @@ def test_call_delete(mocker): mock_delete_routine = mocker.patch( target=f"{path_controller}.Hub.delete_routine", ) + # mock user input + mock_input = "y" + monkeypatch.setattr(f"{path_controller}.console.input", lambda _: mock_input) controller.call_delete( other_args=[ @@ -561,3 +567,85 @@ def test_call_delete(mocker): auth_header="Bearer 123", name="script1", ) + + +def test_call_generate(mocker, monkeypatch): + controller = account_controller.AccountController(queue=None) + path_controller = "openbb_terminal.account.account_controller" + + response = Response() + response.status_code = 200 + response._content = json.dumps( # pylint: disable=protected-access + {"token": "MOCK_TOKEN"} + ).encode("utf-8") + + mocker.patch( + target=f"{path_controller}.User.get_auth_header", + return_value="Bearer 123", + ) + mock_generate = mocker.patch( + target=f"{path_controller}.Hub.generate_personal_access_token", + return_value=response, + ) + + # mock user input + mock_input = "y" + monkeypatch.setattr(f"{path_controller}.console.input", lambda _: mock_input) + + # mock save to keys + mocker.patch( + target=f"{path_controller}.keys_model.set_openbb_personal_access_token", + return_value=True, + ) + + controller.call_generate(other_args=["--save", "--days", "30"]) + + mock_generate.assert_called_once_with( + auth_header="Bearer 123", + days=30, + ) + + +def test_call_show(mocker): + controller = account_controller.AccountController(queue=None) + path_controller = "openbb_terminal.account.account_controller" + + response = Response() + response.status_code = 200 + response._content = json.dumps( # pylint: disable=protected-access + {"token": "MOCK_TOKEN"} + ).encode("utf-8") + + mocker.patch( + target=f"{path_controller}.User.get_auth_header", + return_value="Bearer 123", + ) + mock_get_token = mocker.patch( + target=f"{path_controller}.Hub.get_personal_access_token", + return_value=response, + ) + controller.call_show(other_args=[]) + mock_get_token.assert_called_once_with(auth_header="Bearer 123") + + +def test_call_revoke(mocker, monkeypatch): + controller = account_controller.AccountController(queue=None) + path_controller = "openbb_terminal.account.account_controller" + + response = Response() + response.status_code = 200 + + mocker.patch( + target=f"{path_controller}.User.get_auth_header", + return_value="Bearer 123", + ) + mock_revoke_token = mocker.patch( + target=f"{path_controller}.Hub.revoke_personal_access_token", + return_value=response, + ) + # mock user input + mock_input = "y" + monkeypatch.setattr(f"{path_controller}.console.input", lambda _: mock_input) + + controller.call_revoke(other_args=[]) + mock_revoke_token.assert_called_once_with(auth_header="Bearer 123") diff --git a/tests/openbb_terminal/account/txt/test_account_controller/test_print_help.txt b/tests/openbb_terminal/account/txt/test_account_controller/test_print_help.txt index 01858e13064..0443a8c393f 100644 --- a/tests/openbb_terminal/account/txt/test_account_controller/test_print_help.txt +++ b/tests/openbb_terminal/account/txt/test_account_controller/test_print_help.txt @@ -9,6 +9,11 @@ Cloud storage of routines: download Download routine delete Delete routine +OpenBB Personal Access Token: + generate Generate a Personal Access Token + show Shows the current Personal Access Token + revoke Revoke the Personal Access Token + Authentication: logout Log out from OpenBB account diff --git a/tests/openbb_terminal/session/test_hub_model.py b/tests/openbb_terminal/session/test_hub_model.py index 1aad4b7363a..803503ace75 100644 --- a/tests/openbb_terminal/session/test_hub_model.py +++ b/tests/openbb_terminal/session/test_hub_model.py @@ -1,3 +1,4 @@ +import json from unittest.mock import MagicMock, patch import pytest @@ -709,3 +710,138 @@ def test_list_routines_error(side_effect): ): result = hub_model.list_routines(auth_header="Bearer 123", page=1, size=10) assert result is None + + +@pytest.mark.parametrize( + "auth_header, base_url, timeout, days, status_code", + [ + ("auth_header", "base_url", 10, 10, 200), + ("other_header", "other_url", 10, 10, 400), + ], +) +def test_generate_personal_access_token( + auth_header, base_url, timeout, days, status_code +): + mock_response = MagicMock(spec=requests.Response) + mock_response.status_code = status_code + + with patch( + "openbb_terminal.session.hub_model.requests.put", return_value=mock_response + ) as requests_put_mock: + result = hub_model.generate_personal_access_token( + auth_header=auth_header, base_url=base_url, timeout=timeout, days=days + ) + + assert result.status_code == mock_response.status_code + requests_put_mock.assert_called_once() + _, kwargs = requests_put_mock.call_args + assert kwargs["url"] == base_url + "/sdk/token" + assert kwargs["headers"] == { + "Authorization": auth_header, + "Content-Type": "application/json", + } + assert kwargs["data"] == json.dumps({"days": days}) + assert kwargs["timeout"] == timeout + + +@pytest.mark.parametrize( + "side_effect", + [ + requests.exceptions.ConnectionError, + requests.exceptions.Timeout, + Exception, + ], +) +def test_generate_personal_access_token_error(side_effect): + with patch( + "openbb_terminal.session.hub_model.requests.put", + side_effect=side_effect, + ): + result = hub_model.generate_personal_access_token("auth_header", 10) + assert result is None + + +@pytest.mark.parametrize( + "auth_header, base_url, timeout, status_code", + [ + ("auth_header", "base_url", 10, 200), + ("other_header", "other_url", 10, 400), + ], +) +def test_get_personal_access_token(auth_header, base_url, timeout, status_code): + mock_response = MagicMock(spec=requests.Response) + mock_response.status_code = status_code + + with patch( + "openbb_terminal.session.hub_model.requests.get", return_value=mock_response + ) as requests_get_mock: + result = hub_model.get_personal_access_token( + auth_header=auth_header, base_url=base_url, timeout=timeout + ) + + assert result.status_code == mock_response.status_code + requests_get_mock.assert_called_once() + _, kwargs = requests_get_mock.call_args + assert kwargs["url"] == base_url + "/sdk/token" + assert kwargs["headers"] == {"Authorization": auth_header} + assert kwargs["timeout"] == timeout + + +@pytest.mark.parametrize( + "side_effect", + [ + requests.exceptions.ConnectionError, + requests.exceptions.Timeout, + Exception, + ], +) +def test_get_personal_access_token_error(side_effect): + with patch( + "openbb_terminal.session.hub_model.requests.get", + side_effect=side_effect, + ): + result = hub_model.get_personal_access_token("auth_header") + assert result is None + + +@pytest.mark.parametrize( + "auth_header, base_url, timeout, status_code", + [ + ("auth_header", "base_url", 10, 200), + ("other_header", "other_url", 10, 400), + ], +) +def test_revoke_personal_access_token(auth_header, base_url, timeout, status_code): + mock_response = MagicMock(spec=requests.Response) + mock_response.status_code = status_code + + with patch( + "openbb_terminal.session.hub_model.requests.get", return_value=mock_response + ) as requests_get_mock: + result = hub_model.get_personal_access_token( + auth_header=auth_header, base_url=base_url, timeout=timeout + ) + + assert result.status_code == mock_response.status_code + requests_get_mock.assert_called_once() + _, kwargs = requests_get_mock.call_args + assert kwargs["url"] == base_url + "/sdk/token" + assert kwargs["headers"] == {"Authorization": auth_header} + assert kwargs["timeout"] == timeout + + +@pytest.mark.parametrize( + "side_effect", + [ + requests.exceptions.ConnectionError, + requests.exceptions.Timeout, + Exception, + ], +) +def test_revoke_personal_access_token_error(side_effect): + with patch( + "openbb_terminal.session.hub_model.requests.delete", + side_effect=side_effect, + ): + result = hub_model.revoke_personal_access_token("auth_header") + assert result is None |