diff --git a/openbb_terminal/account/login_controller.py b/openbb_terminal/account/login_controller.py index 5a6e1af78bc..ecd2a1338f1 100644 --- a/openbb_terminal/account/login_controller.py +++ b/openbb_terminal/account/login_controller.py @@ -13,7 +13,7 @@ from openbb_terminal.account.login_model import ( def display_welcome_message(): """Display welcome message""" console.print( - "You need to log in to use the installer version of OpenBB.", + "\nYou must log in to use the installer version of OpenBB.", style="info", ) diff --git a/openbb_terminal/account/login_model.py b/openbb_terminal/account/login_model.py index cc4c32ac157..1a26310ed00 100644 --- a/openbb_terminal/account/login_model.py +++ b/openbb_terminal/account/login_model.py @@ -1,4 +1,3 @@ -from dataclasses import dataclass from enum import Enum from pathlib import Path import os.path @@ -11,12 +10,15 @@ from openbb_terminal.core.config.paths import SETTINGS_DIRECTORY from openbb_terminal.rich_config import console from openbb_terminal import feature_flags +from openbb_terminal.account.user import User + BASE_URL = "http://127.0.0.1:8000/terminal/" class Status(Enum): VALID_LOGIN = "[green]\nLogin successful.[/green]" INVALID_LOGIN = "[red]\nInvalid login.[/red]" + LOGGED_OUT = "[green]\nLogged out.[/green]" class Error(Enum): @@ -28,13 +30,6 @@ class Error(Enum): READ_ERROR = "[red]\nError reading login info.[/red]" -# Will something like this for the logging -@dataclass(frozen=True) -class User: - email: str - uuid: str - - def launch_terminal(login_info: dict): """Launch the terminal. @@ -43,13 +38,16 @@ def launch_terminal(login_info: dict): login_info : dict The login info. """ - token = login_info.get("access_token", "") + User.token_type = login_info.get("token_type", "") + User.token = login_info.get("access_token", "") User.uuid = login_info.get("uuid", "") - if token: - decoded_info = jwt.decode(token, options={"verify_signature": False}) + if User.token: + decoded_info = jwt.decode(User.token, options={"verify_signature": False}) User.email = decoded_info.get("sub", "") - username = User.email[: User.email.find("@")] - setattr(feature_flags, "USE_FLAIR", "[" + username + "] 🦋") + + if feature_flags.USE_FLAIR == ":openbb": + username = User.email[: User.email.find("@")] + setattr(feature_flags, "USE_FLAIR", "[" + username + "] 🦋") terminal_controller.parse_args_and_run() @@ -65,8 +63,8 @@ def write_to_file(data: dict, file_path: Path): The file path. """ try: - with open(file_path, "w") as outfile: - outfile.write(json.dumps(data)) + with open(file_path, "w") as file: + file.write(json.dumps(data)) except Exception: console.print(Error.SAVE_ERROR.value) @@ -130,7 +128,9 @@ def request_login_info(email: str, password: str, save: bool) -> Union[dict, Err except requests.exceptions.ConnectionError: console.print(Error.CONNECTION_ERROR.value) return Error.CONNECTION_ERROR - + except Exception: + console.print(Error.UNKNOWN_ERROR.value) + return Error.UNKNOWN_ERROR return process_response(response, save) @@ -183,4 +183,34 @@ def get_login_status(login_info: dict) -> Union[Status, Error]: except requests.exceptions.ConnectionError: console.print(Error.CONNECTION_ERROR.value) return Error.CONNECTION_ERROR + except Exception: + console.print(Error.UNKNOWN_ERROR.value) + return Error.UNKNOWN_ERROR return Status.INVALID_LOGIN + + +def logout() -> Union[Status, Error]: + """Logout the user. + + Returns + ------- + Union[Status, Error] + The status of the logout, or an error. + """ + try: + # TODO: Uncomment when the server supports logout + # requests.post( + # url=BASE_URL + "logout", + # headers={"Authorization": f"{User.token_type.title()} {User.token}"}, + # ) + file_path = SETTINGS_DIRECTORY / "login.json" + if os.path.isfile(file_path): + os.remove(file_path) + console.print(Status.LOGGED_OUT.value) + return Status.LOGGED_OUT + except requests.exceptions.ConnectionError: + console.print(Error.CONNECTION_ERROR.value) + return Error.CONNECTION_ERROR + except Exception: + console.print(Error.UNKNOWN_ERROR.value) + return Error.UNKNOWN_ERROR diff --git a/openbb_terminal/account/user.py b/openbb_terminal/account/user.py new file mode 100644 index 00000000000..b874abc053a --- /dev/null +++ b/openbb_terminal/account/user.py @@ -0,0 +1,10 @@ +# Will something like this for the logging +from dataclasses import dataclass + + +@dataclass(frozen=True) +class User: + token_type: str = "" + token: str = "" + email: str = "" + uuid: str = "" diff --git a/openbb_terminal/parent_classes.py b/openbb_terminal/parent_classes.py index c719edccd7f..87c5a16f706 100644 --- a/openbb_terminal/parent_classes.py +++ b/openbb_terminal/parent_classes.py @@ -18,6 +18,8 @@ import pandas as pd from prompt_toolkit.formatted_text import HTML from prompt_toolkit.styles import Style from rich.markdown import Markdown +from openbb_terminal.account.login_model import Status, logout +from openbb_terminal.account.user import User from openbb_terminal.core.config.paths import ( USER_CUSTOM_IMPORTS_DIRECTORY, @@ -47,6 +49,7 @@ from openbb_terminal.stocks import stocks_helper from openbb_terminal.terminal_helper import open_openbb_documentation from openbb_terminal.cryptocurrency import cryptocurrency_helpers from openbb_terminal.core.completer.choices import build_controller_choice_map +from openbb_terminal.account import login_controller logger = logging.getLogger(__name__) @@ -94,6 +97,8 @@ class BaseController(metaclass=ABCMeta): "record", "stop", "screenshot", + "logout", + "whoami", ] CHOICES_COMMANDS: List[str] = [] @@ -656,6 +661,36 @@ class BaseController(metaclass=ABCMeta): if ns_parser: screenshot() + @log_start_end(log=logger) + def call_logout(self, other_args: List[str]) -> None: + """Process logout command.""" + parser = argparse.ArgumentParser( + add_help=False, + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + prog="logout", + description="Logout from OpenBB", + ) + ns_parser = self.parse_simple_args(parser, other_args) + + if ns_parser: + if logout() == Status.LOGGED_OUT: + login_controller.main() + + @log_start_end(log=logger) + def call_whoami(self, other_args: List[str]) -> None: + """Process whoami command.""" + parser = argparse.ArgumentParser( + add_help=False, + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + prog="whoami", + description="Show current user", + ) + ns_parser = self.parse_simple_args(parser, other_args) + + if ns_parser: + console.print(f"[info]email:[/info] {User.email}") + console.print(f"[info]uuid:[/info] {User.uuid}\n") + @staticmethod def parse_simple_args(parser: argparse.ArgumentParser, other_args: List[str]): """Parse list of arguments into the supplied parser.