CDRM-Project/routes/api.py

"""Module to handle the API routes."""

import os
import sqlite3
import json
import shutil
import math
from io import StringIO
import tempfile
import time

from flask import Blueprint, jsonify, request, send_file, session, after_this_request
import yaml
import mysql.connector

from custom_functions.decrypt.api_decrypt import api_decrypt
from custom_functions.user_checks.device_allowed import user_allowed_to_use_device
from custom_functions.database.unified_db_ops import (
    search_by_pssh_or_kid,
    cache_to_db,
    get_key_by_kid_and_service,
    get_unique_services,
    get_kid_key_dict,
    key_count,
)
from configs.icon_links import data as icon_data

api_bp = Blueprint("api", __name__)
with open(os.path.join(os.getcwd(), "configs", "config.yaml"), "r", encoding="utf-8") as file:
    config = yaml.safe_load(file)


def get_db_config():
    """Get the MariaDB database configuration."""
    with open(
        os.path.join(os.getcwd(), "configs", "config.yaml"), "r", encoding="utf-8"
    ) as file_mariadb:
        config_mariadb = yaml.safe_load(file_mariadb)
    db_config = {
        "host": f'{config_mariadb["mariadb"]["host"]}',
        "user": f'{config_mariadb["mariadb"]["user"]}',
        "password": f'{config_mariadb["mariadb"]["password"]}',
        "database": f'{config_mariadb["mariadb"]["database"]}',
    }
    return db_config


@api_bp.route("/api/cache/search", methods=["POST"])
def get_data():
    """Get the data from the database."""
    search_argument = json.loads(request.data)["input"]
    results = search_by_pssh_or_kid(search_filter=search_argument)
    return jsonify(results)


@api_bp.route("/api/cache/<service>/<kid>", methods=["GET"])
def get_single_key_service(service, kid):
    """Get the single key from the database."""
    result = get_key_by_kid_and_service(kid=kid, service=service)
    return jsonify(
        {
            "code": 0,
            "content_key": result,
        }
    )


@api_bp.route("/api/cache/<service>", methods=["GET"])
def get_multiple_key_service(service):
    """Get the multiple keys from the database."""
    result = get_kid_key_dict(service_name=service)
    pages = math.ceil(len(result) / 10)
    return jsonify({"code": 0, "content_keys": result, "pages": pages})


@api_bp.route("/api/cache/<service>/<kid>", methods=["POST"])
def add_single_key_service(service, kid):
    """Add the single key to the database."""
    body = request.get_json()
    content_key = body["content_key"]
    result = cache_to_db(service=service, kid=kid, key=content_key)
    if result:
        return jsonify(
            {
                "code": 0,
                "updated": True,
            }
        )
    return jsonify(
        {
            "code": 0,
            "updated": True,
        }
    )


@api_bp.route("/api/cache/<service>", methods=["POST"])
def add_multiple_key_service(service):
    """Add the multiple keys to the database."""
    body = request.get_json()
    keys_added = 0
    keys_updated = 0
    for kid, key in body["content_keys"].items():
        result = cache_to_db(service=service, kid=kid, key=key)
        if result is True:
            keys_updated += 1
        else:
            keys_added += 1
    return jsonify(
        {
            "code": 0,
            "added": str(keys_added),
            "updated": str(keys_updated),
        }
    )


@api_bp.route("/api/cache", methods=["POST"])
def unique_service():
    """Get the unique services from the database."""
    services = get_unique_services()
    return jsonify(
        {
            "code": 0,
            "service_list": services,
        }
    )


@api_bp.route("/api/cache/download", methods=["GET"])
def download_database():
    """Download the database."""
    if config["database_type"].lower() != "mariadb":
        original_database_path = os.path.join(os.getcwd(), "databases", "sql", "key_cache.db")

        # Make a copy of the original database (without locking the original)
        modified_database_path = os.path.join(os.getcwd(), "databases", "sql", "key_cache_modified.db")

        # Using shutil.copy2 to preserve metadata (timestamps, etc.)
        shutil.copy2(original_database_path, modified_database_path)

        # Open the copied database for modification using 'with' statement to avoid locks
        with sqlite3.connect(modified_database_path) as conn:
            cursor = conn.cursor()

            # Update all rows to remove Headers and Cookies (set them to NULL or empty strings)
            cursor.execute(
                """
            UPDATE licenses
            SET Headers = NULL,
                Cookies = NULL
            """
            )

            # No need for explicit commit, it's done automatically with the 'with' block
            # The connection will automatically be committed and closed when the block ends

        # Send the modified database as an attachment
        return send_file(
            modified_database_path, as_attachment=True, download_name="key_cache.db"
        )
    try:
        conn = mysql.connector.connect(**get_db_config())
        cursor = conn.cursor()

        # Get column names
        cursor.execute("SHOW COLUMNS FROM licenses")
        columns = [row[0] for row in cursor.fetchall()]

        # Build SELECT with Headers and Cookies as NULL
        select_columns = []
        for col in columns:
            if col.lower() in ("headers", "cookies"):
                select_columns.append("NULL AS " + col)
            else:
                select_columns.append(col)
        select_query = f"SELECT {', '.join(select_columns)} FROM licenses"
        cursor.execute(select_query)
        rows = cursor.fetchall()

        # Dump to SQL-like format
        output = StringIO()
        output.write("-- Dump of `licenses` table (Headers and Cookies are NULL)\n")
        for row in rows:
            values = ", ".join(
                f"'{str(v).replace('\'', '\\\'')}'" if v is not None else "NULL"
                for v in row
            )
            output.write(
                f"INSERT INTO licenses ({', '.join(columns)}) VALUES ({values});\n"
            )

        # Write to a temp file for download
        temp_dir = tempfile.gettempdir()
        temp_path = os.path.join(temp_dir, "key_cache.sql")
        with open(temp_path, "w", encoding="utf-8") as f:
            f.write(output.getvalue())

        @after_this_request
        def remove_file(response):
            try:
                os.remove(temp_path)
            except Exception:
                pass
            return response

        return send_file(
            temp_path, as_attachment=True, download_name="licenses_dump.sql"
        )
    except mysql.connector.Error as err:
        return {"error": str(err)}, 500


_keycount_cache = {"count": None, "timestamp": 0}


@api_bp.route("/api/cache/keycount", methods=["GET"])
def get_count():
    """Get the count of the keys in the database."""
    now = time.time()
    if now - _keycount_cache["timestamp"] > 10 or _keycount_cache["count"] is None:
        _keycount_cache["count"] = key_count()
        _keycount_cache["timestamp"] = now
    return jsonify({"count": _keycount_cache["count"]})


@api_bp.route("/api/decrypt", methods=["POST"])
def decrypt_data():
    """Decrypt the data."""
    api_request_data = request.get_json(force=True)

    # Helper to get fields or None if missing/empty
    def get_field(key, default=""):
        value = api_request_data.get(key, default)
        return value if value != "" else default

    api_request_pssh = get_field("pssh")
    api_request_licurl = get_field("licurl")
    api_request_proxy = get_field("proxy")
    api_request_headers = get_field("headers")
    api_request_cookies = get_field("cookies")
    api_request_data_func = get_field("data")

    # Device logic
    device = get_field("device", "public")
    if device in [
        "default",
        "CDRM-Project Public Widevine CDM",
        "CDRM-Project Public PlayReady CDM",
        "",
        None,
    ]:
        api_request_device = "public"
    else:
        api_request_device = device

    username = ""
    if api_request_device != "public":
        username = session.get("username")
        if not username:
            return jsonify({"message": "Not logged in, not allowed"}), 400
        if not user_allowed_to_use_device(device=api_request_device, username=username):
            return jsonify({"message": "Not authorized / Not found"}), 403

    result = api_decrypt(
        pssh=api_request_pssh,
        proxy=api_request_proxy,
        license_url=api_request_licurl,
        headers=api_request_headers,
        cookies=api_request_cookies,
        json_data=api_request_data_func,
        device=api_request_device,
        username=username,
    )
    if result["status"] == "success":
        return jsonify({"status": "success", "message": result["message"]})
    return jsonify({"status": "fail", "message": result["message"]})


@api_bp.route("/api/links", methods=["GET"])
def get_links():
    """Get the links."""
    return jsonify(
        {
            "discord": icon_data["discord"],
            "telegram": icon_data["telegram"],
            "gitea": icon_data["gitea"],
        }
    )


@api_bp.route("/api/extension", methods=["POST"])
def verify_extension():
    """Verify the extension."""
    return jsonify(
        {
            "status": True,
        }
    )