yt-dlp-host

1
from functools import wraps
2
from flask import request, jsonify
3
from src.json_utils import load_keys, save_keys, load_tasks
4
from config import REQUEST_LIMIT, TASK_CLEANUP_TIME
5
import secrets
6

7
def generate_key():
8
    return secrets.token_urlsafe(32)
9

10
def check_rate_limit(api_key):
11
    tasks = load_tasks()
12
    key_name = get_key_name(api_key)
13
    rate = 0
14
    for task_name, task_info in tasks.items():
15
        if task_info['key_name'] == key_name:
16
            rate += 1
17
    if rate >= REQUEST_LIMIT: return False
18
    return True
19

20
def check_api_key(required_permission):
21
    def decorator(f):
22
        @wraps(f)
23
        def decorated_function(*args, **kwargs):
24
            api_key = request.headers.get('X-API-Key')
25
            if not api_key:
26
                return jsonify({'error': 'No API key provided'}), 401
27
            if not check_rate_limit(api_key):
28
                return jsonify({'error': f'Rate limit exceeded. Maximum {REQUEST_LIMIT} requests per {TASK_CLEANUP_TIME} minutes.'}), 429
29
            key_info = get_key_info(api_key)
30
            if not key_info:
31
                return jsonify({'error': 'Invalid API key'}), 401
32
            permissions = key_info['permissions']
33
            if required_permission not in permissions:
34
                return jsonify({'error': 'Insufficient permissions'}), 403
35
            return f(*args, **kwargs)
36
        return decorated_function
37
    return decorator
38

39
def get_key_name(api_key):
40
    keys = load_keys()
41
    for key_name, key_info in keys.items():
42
        if key_info['key'] == api_key:
43
            return key_name
44
    return None
45

46
def create_api_key(name, permissions):
47
    keys = load_keys()
48
    new_key = generate_key()
49
    keys[name] = {
50
        'key': new_key,
51
        'permissions': permissions
52
    }
53
    save_keys(keys)
54
    return new_key
55

56
def delete_api_key(name):
57
    keys = load_keys()
58
    if name in keys:
59
        del keys[name]
60
        save_keys(keys)
61
        return True
62
    return False
63

64
def get_all_keys():
65
    return load_keys()
66

67
def get_key_info(api_key):
68
    keys = load_keys()
69
    key_info = next((item for item in keys.values() if item['key'] == api_key), None)
70
    return key_info
71

72
if load_keys() == {}: create_api_key("admin", ["create_key", "delete_key", "get_key", "get_keys", "get_video", "get_audio", "get_info"])
73