OpenBB/openbb_terminal/core/session/hub_model.py

import json
from typing import Dict, Literal, Optional

import requests

from openbb_terminal.core.session.constants import (
    BASE_URL,
    CONNECTION_ERROR_MSG,
    CONNECTION_TIMEOUT_MSG,
    TIMEOUT,
)
from openbb_terminal.core.session.current_system import get_current_system
from openbb_terminal.rich_config import console


def create_session(
    email: str, password: str, base_url: str = BASE_URL, timeout: int = TIMEOUT
) -> Optional[requests.Response]:
    """Create a session.

    Parameters
    ----------
    email : str
        The email.
    password : str
        The password.
    base_url : str
        The base url, by default BASE_URL
    timeout : int
        The timeout, by default TIMEOUT

    Returns
    -------
    Optional[requests.Response]
        The response from the login request.
    """
    try:
        data = {
            "email": email,
            "password": password,
            "remember": True,
        }
        return requests.post(url=base_url + "login", json=data, timeout=timeout)
    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("\n[red]Failed to request login info.[/red]")
        return None


def create_session_from_token(
    token: str, base_url: str = BASE_URL, timeout: int = TIMEOUT
):
    """Create a session from token.

    Parameters
    ----------
    token : str
        The token.
    base_url : str
        The base url, by default BASE_URL
    timeout : int
        The timeout, by default TIMEOUT
    """
    try:
        data = {
            "token": token,
        }
        return requests.post(url=base_url + "sdk/login", json=data, timeout=timeout)
    except requests.exceptions.ConnectionError:
        console.print("\n[red]Connection error.[/red]")
        return None
    except requests.exceptions.Timeout:
        console.print("\n[red]Connection timeout.[/red]")
        return None
    except Exception:
        console.print("\n[red]Failed to request login info.[/red]")
        return None


def delete_session(
    auth_header: str, token: str, base_url: str = BASE_URL, timeout: int = TIMEOUT
) -> Optional[requests.Response]:
    """Delete the session.

    Parameters
    ----------
    auth_header : str
        The authorization header, e.g. "Bearer <token>".
    token : str
        The token to delete.
    base_url : str
        The base url, by default BASE_URL
    timeout : int
        The timeout, by default TIMEOUT

    Returns
    -------
    Optional[requests.Response]
        The response from the logout request.
    """
    try:
        response = requests.get(
            url=base_url + "logout",
            headers={"Authorization": auth_header},
            json={"token": token},
            timeout=timeout,
        )
        if response.status_code != 200:
            console.print("[red]Failed to delete server session.[/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 delete server session.[/red]")
        return None


def process_session_response(response: requests.Response) -> dict:
    """Process the response from the login request.

    Parameters
    ----------
    response : requests.Response
        The response from the login request.

    Returns
    -------
    dict
        The login info.
    """
    if response.status_code == 200:
        login = response.json()
        return login
    if response.status_code == 401:
        console.print("\n[red]Wrong credentials.[/red]")
        return {}
    if response.status_code == 403:
        console.print("\n[red]Unverified email.[/red]")
        return {}
    console.print("\n[red]Failed to login.[/red]")
    return {}


def get_session(email: str, password: str) -> dict:
    """Get the session info.

    Parameters
    ----------
    email : str
        The email.
    password : str
        The password.

    Returns
    -------
    dict
        The session info.
    """
    response = create_session(email, password)
    if response is None:
        return {}
    return process_session_response(response)


def get_session_from_token(token: str) -> dict:
    """Get the session info from token.

    Parameters
    ----------
    token : str
        The token.

    Returns
    -------
    dict
        The session info.
    """
    response = create_session_from_token(token)
    if response is None:
        return {}
    return process_session_response(response)


def fetch_user_configs(
    session: dict, base_url: str = BASE_URL, timeout: int = TIMEOUT
) -> Optional[requests.Response]:
    """Fetch user configurations.

    Parameters
    ----------
    session : dict
        The session info.
    base_url : str
        The base url, by default BASE_URL
    timeout : int
        The timeout, by default TIMEOUT

    Returns
    -------
    Optional[requests.Response]
        The response from the get request.
    """
    token_type = session.get("token_type", "")
    token = session.get("access_token", "")

    try:
        response = requests.get(
            url=base_url + "terminal/user",
            headers={"Authorization": f"{token_type.title()} {token}"},
            timeout=timeout,
        )
        if response.status_code not in [200, 401]:  # 401: invalid token
            console.print("[red]\nFailed to fetch configurations.[/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]\nFailed to fetch configurations.[/red]")
        return None


def upload_config(
    key: str,
    value: str,
    type_: Literal["keys", "settings", "sources"],
    auth_header: str,
    base_url: str = BASE_URL,
    timeout: int = TIMEOUT,
) -> Optional[requests.Response]:
    """Patch user configurations to the server.

    Parameters
    ----------
    key : str
        The key to patch.
    value : str
        The value to patch.
    type_ : Literal["keys", "settings", "sources"]
        The type of the patch.
    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]
        The response from the patch request.
    """
    if type_ not in ["keys", "settings", "sources"]:
        console.print("[red]\nInvalid patch type.[/red]")
        return None

    data = {"key": f"features_{type_}.{key}", "value": value}

    try:
        response = requests.patch(
            url=base_url + "terminal/user",
            headers={"Authorization": auth_header},
            json=data,
            timeout=timeout,
        )
        if response.status_code == 200:
            console.print("[green]Saved remotely.[/green]")
        else:
            console.print("[red]Failed to save remotely.[/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 save remotely.[/red]")
        return None


def clear_user_configs(
    auth_header: str, base_url: str = BASE_URL, timeout: int = TIMEOUT
) -> Optional[requests.Response]:
    """Clear user configurations to the server.

    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]
        The response from the put request.
    """
    data: Dict[str, dict] = {"features_keys": {}}

    try:
        response = requests.put(
            url=base_url + "user",
            headers={"Authorization": auth_header},
            json=data,
            timeout=timeout,
        )
        if response.status_code == 200:
            console.print("[green]Cleared data.[/green]")
        else:
            console.print("[red]Failed to clear data.[/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 clear data.[/red]")
        return None


def upload_routine(
    auth_header: str,
    name: str = "",
    description: str = "",
    routine: str = "",
    override: bool = False,
    base_url=BASE_URL,
    timeout: int = TIMEOUT,
) -> Optional[requests.Response]:
    """Send a routine to the server.

    Parameters
    ----------
    auth_header : str
        The authorization header, e.g. "Bearer <token>".
    name : str
        The name of the routine.
    routine : str
        The routine.
    override : bool
        Whether to override the routine if it already exists.
    base_url : str
        The base url, by default BASE_URL
    timeout : int
        The timeout, by default TIMEOUT

    Returns
    -------
    Optional[requests.Response]
        The response from the post request.
    """
    data = {
        "name": name,
        "description": description,
        "script": routine,
        "override": override,
        "version": get_current_system().VERSION,
        "public": False,
    }

    try:
        response = requests.post(
            headers={"Authorization": auth_header},
            url=base_url + "terminal/script",
            json=data,
            timeout=timeout,
        )
        if response.status_code == 200:
            console.print("[green]Successfully uploaded your routine.[/green]")
        elif response.status_code != 409:  # 409: routine already exists
            console.print(
                "[red]" + response.json().get("detail", "Unknown error.") + "[/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 upload your routine.[/red]")
        return None


def download_routine(
    auth_header: str,
    name: str = "",
    base_url=BASE_URL,
    timeout: int = TIMEOUT,
) -> Optional[requests.Response]:
    """Download a routine from the server.

    Parameters
    ----------
    auth_header : str
        The authorization header, e.g. "Bearer <token>".
    name : str
        The name of the routine.
    base_url : str
        The base url, by default BASE_URL
    timeout : int
        The timeout, by default TIMEOUT

    Returns
    -------
    Optional[requests.Response]
        The response from the get request.
    """
    try:
        response = requests.get(
            headers={"Authorization": auth_header},
            url=base_url + "terminal/script/" + name,
            timeout=timeout,
        )
        if response.status_code == 404:
            console.print("[red]Routine not found.[/red]")
        elif response.status_code != 200:
            console.print("[red]Failed to download your routine.[/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 download your routine.[/red]")
        return None


def delete_routine(
    auth_header: str,
    name: str = "",
    base_url=BASE_URL,
    timeout: int = TIMEOUT,
) -> Optional[requests.Response]:
    """Delete a routine from the server.

    Parameters
    ----------
    auth_header : str
        The authorization header, e.g. "Bearer <token>".
    name : str
        The name of the routine.
    base_url : str
        The base url, by default BASE_URL
    timeout : int
        The timeout, by default TIMEOUT

    Returns
    -------
    Optional[requests.Response]
        The response from the delete request.
    """
    try:
        response = requests.delete(
            headers={"Authorization": auth_header},
            url=base_url + "terminal/script/" + name,
            timeout=timeout,
        )
        if response.status_code == 200:
            console.print("[green]Successfully deleted your routine.[/green]")
        elif response.status_code == 404:
            console.print("[red]Routine not found.[/red]")
        else:
            console.print("[red]Failed to delete your routine.[/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 delete your routine.[/red]")
        return None


def list_routines(
    auth_header: str,
    page: int = 1,
    size: int = 10,
    base_url=BASE_URL,
    timeout: int = TIMEOUT,
) -> Optional[requests.Response]:
    """List all routines from the server.

    Parameters
    ----------
    auth_header : str
        The authorization header, e.g. "Bearer <token>".
    page : int
        The page number.
    size : int
        The number of routines per page.
    base_url : str
        The base url, by default BASE_URL
    timeout : int
        The timeout, by default TIMEOUT

    Returns
    -------
    Optional[requests.Response]
        The response from the get request.
    """
    fields = "name%2Cdescription"

    try:
        response = requests.get(
            headers={"Authorization": auth_header},
            url=f"{base_url}terminal/script?fields={fields}&page={page}&size={size}",
            timeout=timeout,
        )
        if response.status_code != 200:
            console.print("[red]Failed to list your routines.[/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 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