Source code for psnawp_api.utils.misc

"""Miscellaneous functions that are used throughout the psnawp module."""

from __future__ import annotations

import base64
import binascii
import json
import os
import re
from datetime import datetime
from pathlib import Path
from tempfile import gettempdir
from time import time
from typing import TYPE_CHECKING, cast

from pycountry import countries

from psnawp_api.core.psnawp_exceptions import PSNAWPInvalidTokenError

if TYPE_CHECKING:
    from pycountry.db import Country

ISO_3166_1_ALPHA_2_LEN = 2




def iso_format_to_datetime(iso_format: str | None) -> datetime | None:
    """Converts an ISO 8601 formatted string to a :py:class:`~datetime.datetime` object.

    :param iso_format: The ISO 8601 formatted string (e.g., "2025-02-26T12:00:00Z").

    :returns: The corresponding :py:class:`~datetime.datetime` object, or ``None`` if input is ``None``.

    """
    return datetime.fromisoformat(iso_format.replace("Z", "+00:00")) if iso_format is not None else None





def extract_region_from_npid(npid: str) -> Country | None:
    """Extracts the region code from a base64-encoded NPID string and converts it to a full country.

    The function decodes the NPID, extracts the last two characters representing the ISO 3166-1 alpha-2 region code, and
    uses the `pycountry` library to map it to a full country.

    :param npid: The base64-encoded NPID string, which is decoded to extract the region code.

    :returns: The region as a :py:class:`~pycountry.db.Country` object, or ``None`` if extraction fails or region is
        invalid.

    This function assumes a valid NPID format (e.g., `"VaultTec-Co@b7.us"`). It splits the string to get the region code
    (e.g., "US"), which is then matched to the full country using the `pycountry` library.

    """
    try:
        decoded_npid = base64.b64decode(npid).decode("utf-8")
    except (binascii.Error, UnicodeDecodeError):
        return None

    # Assuming a valid decoded npid format (e.g. VaultTec-Co@b7.us), extract the region (e.g. "US")
    if "@" in decoded_npid and "." in decoded_npid:
        region_candidate = decoded_npid.split(".")[-1]
        if len(region_candidate) == ISO_3166_1_ALPHA_2_LEN and region_candidate.isalpha():
            return cast("Country", countries.get(alpha_2=region_candidate))

    return None





def parse_npsso_token(npsso_input: str) -> str:
    """Accept string from the user that may contain either a valid npsso token or a json string with key "npsso" and value of the npsso token.

    This function either succeeds at extracting the npsso token from the provided input (meaning a valid npsso json
    string was provided) or it returns the original input.

    :param npsso_input: User provided input for npsso token.

    :returns: Extracted npsso token from user input or the original string.

    :raises PSNAWPInvalidTokenError: If malformed npsso JSON is supplied

    """
    pattern = r"\{|\}"
    if re.search(pattern, npsso_input):
        try:
            npsso_dict: dict[str, str] = json.loads(npsso_input)
            return npsso_dict["npsso"]
        except json.JSONDecodeError as exp:
            raise PSNAWPInvalidTokenError("Malformed JSON passed as input.") from exp
        except KeyError as exp:
            raise PSNAWPInvalidTokenError('Input JSON is missing the "npsso" key') from exp
    return npsso_input





def get_temp_db_path(filename: str | None = None) -> Path:
    """Create a writable temporary directory and database file.

    :param filename: Name of the SQLite database file.

    :returns: Path to the writable SQLite database.

    """
    if filename is None:
        filename = f"psnawp_limiter_{time()}.sqlite"

    temp_dir = Path(gettempdir()) / f"psnawp-{os.getpid()}"
    temp_dir.mkdir(parents=True, exist_ok=True)

    db_path = temp_dir / filename
    db_path.touch(exist_ok=True)

    return db_path