CDRM-Project/routes/api.py

import os
import sqlite3
from flask import Blueprint, jsonify, request, send_file, session
import json
from custom_functions.decrypt.api_decrypt import api_decrypt
from custom_functions.user_checks.device_allowed import user_allowed_to_use_device
import shutil
import math
import yaml
import mysql.connector
from io import StringIO
import tempfile
import time
from configs.icon_links import data as icon_data

api_bp = Blueprint('api', __name__)
with open(f'{os.getcwd()}/configs/config.yaml', 'r') as file:
    config = yaml.safe_load(file)
if config['database_type'].lower() != 'mariadb':
    from custom_functions.database.cache_to_db_sqlite import search_by_pssh_or_kid, cache_to_db, \
        get_key_by_kid_and_service, get_unique_services, get_kid_key_dict, key_count
elif config['database_type'].lower() == 'mariadb':
    from custom_functions.database.cache_to_db_mariadb import search_by_pssh_or_kid, cache_to_db, \
        get_key_by_kid_and_service, get_unique_services, get_kid_key_dict, key_count

def get_db_config():
    # Configure your MariaDB connection
    with open(f'{os.getcwd()}/configs/config.yaml', 'r') as file:
        config = yaml.safe_load(file)
    db_config = {
        'host': f'{config["mariadb"]["host"]}',
        'user': f'{config["mariadb"]["user"]}',
        'password': f'{config["mariadb"]["password"]}',
        'database': f'{config["mariadb"]["database"]}'
    }
    return db_config

@api_bp.route('/api/cache/search', methods=['POST'])
def get_data():
    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):
    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):
    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):
    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,
        })
    elif result is False:
        return jsonify({
            'code': 0,
            'updated': True,
        })

@api_bp.route('/api/cache/<service>', methods=['POST'])
def add_multiple_key_service(service):
    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
        elif result is False:
            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():
    services = get_unique_services()
    return jsonify({
        'code': 0,
        'service_list': services,
    })


@api_bp.route('/api/cache/download', methods=['GET'])
def download_database():
    if config['database_type'].lower() != 'mariadb':
        original_database_path = f'{os.getcwd()}/databases/sql/key_cache.db'

        # Make a copy of the original database (without locking the original)
        modified_database_path = f'{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')
    if config['database_type'].lower() == 'mariadb':
        try:
            # Connect to MariaDB
            conn = mysql.connector.connect(**get_db_config())
            cursor = conn.cursor()

            # Update sensitive data (this updates the live DB, you may want to duplicate rows instead)
            cursor.execute('''
            UPDATE licenses
            SET Headers = NULL,
                Cookies = NULL
            ''')

            conn.commit()

            # Now export the table
            cursor.execute('SELECT * FROM licenses')
            rows = cursor.fetchall()
            column_names = [desc[0] for desc in cursor.description]

            # Dump to SQL-like format
            output = StringIO()
            output.write(f"-- Dump of `licenses` table\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(column_names)}) 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())

            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():
    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():
    api_request_data = json.loads(request.data)
    if 'pssh' in api_request_data:
        if api_request_data['pssh'] == '':
            api_request_pssh = None
        else:
            api_request_pssh = api_request_data['pssh']
    else:
        api_request_pssh = None
    if 'licurl' in api_request_data:
        if api_request_data['licurl'] == '':
            api_request_licurl = None
        else:
            api_request_licurl = api_request_data['licurl']
    else:
        api_request_licurl = None
    if 'proxy' in api_request_data:
        if api_request_data['proxy'] == '':
            api_request_proxy = None
        else:
            api_request_proxy = api_request_data['proxy']
    else:
        api_request_proxy = None
    if 'headers' in api_request_data:
        if api_request_data['headers'] == '':
            api_request_headers = None
        else:
            api_request_headers = api_request_data['headers']
    else:
        api_request_headers = None
    if 'cookies' in api_request_data:
        if api_request_data['cookies'] == '':
            api_request_cookies = None
        else:
            api_request_cookies = api_request_data['cookies']
    else:
        api_request_cookies = None
    if 'data' in api_request_data:
        if api_request_data['data'] == '':
            api_request_data_func = None
        else:
            api_request_data_func = api_request_data['data']
    else: api_request_data_func = None
    if 'device' in api_request_data:
        if api_request_data['device'] == 'default' or api_request_data['device'] == 'CDRM-Project Public Widevine CDM' or api_request_data['device'] == 'CDRM-Project Public PlayReady CDM':
            api_request_device = 'public'
        else:
            api_request_device = api_request_data['device']
    else:
        api_request_device = 'public'
    username = None
    if api_request_device != 'public':
        username = session.get('username')
        if not username:
            return jsonify({'message': 'Not logged in, not allowed'}), 400
        if user_allowed_to_use_device(device=api_request_device, username=username):
            api_request_device = api_request_device
        else:
            return jsonify({'message': f'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']
        })
    else:
        return jsonify({
            'status': 'fail',
            'message': result['message']
        })

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