#!/usr/bin/env python3
"""离线下载管理：凭据只在本地读取，API认证只保留在内存。"""
import argparse
import base64
import configparser
import copy
import hashlib
import hmac
import json
import os
from pathlib import Path
import subprocess
import time
import urllib.request
import urllib.error
from hashids import Hashids

BASE = Path('/opt/DockerDate/cloudreve-aria2')
TEMP = Path('/media/ztl/LinuxDate/cloudreve/offline-download')
PRIVATE = Path('/root/AI/cloudreve/private/offline-download')
SITE = 'http://192.168.10.62:5212'


def mysql(sql):
    c = configparser.ConfigParser(interpolation=None)
    c.read('/opt/cloudreve/data/conf.ini')
    db = c['Database']
    options = '[client]\nuser=' + db['User'] + '\npassword=' + db['Password'] + '\n'
    r = subprocess.run(['docker', 'exec', '-i', '1Panel-mysql-5qkq', 'mysql',
        '--defaults-extra-file=/dev/stdin', '--default-character-set=utf8mb4',
        '--protocol=TCP', '-h', '127.0.0.1', '-D', db['Name'], '--batch',
        '--skip-column-names', '-e', sql], input=options, text=True, capture_output=True, timeout=60)
    if r.returncode:
        raise RuntimeError('数据库只读查询失败')
    return r.stdout.rstrip('\n')


def settings():
    result = {}
    for row in mysql("SELECT name,HEX(value) FROM settings WHERE name IN ('secret_key','hash_id_salt')").splitlines():
        name, value = row.split('\t')
        result[name] = bytes.fromhex(value).decode()
    assert result.get('secret_key') and result.get('hash_id_salt'), '认证设置缺失'
    return result


def api(method, path, body=None):
    s = settings()
    enc = lambda x: base64.urlsafe_b64encode(json.dumps(x, separators=(',', ':')).encode()).rstrip(b'=')
    now = int(time.time())
    header = enc({'alg': 'HS256', 'typ': 'JWT'})
    payload = enc({'sub': Hashids(salt=s['hash_id_salt']).encode(1, 1),
                   'token_type': 'access', 'nbf': now - 5, 'exp': now + 120})
    message = header + b'.' + payload
    sig = base64.urlsafe_b64encode(hmac.new(s['secret_key'].encode(), message, hashlib.sha256).digest()).rstrip(b'=')
    token = (message + b'.' + sig).decode()
    req = urllib.request.Request(SITE + '/api/v4' + path, method=method,
        data=json.dumps(body).encode() if body is not None else None,
        headers={'Authorization': 'Bearer ' + token, 'Content-Type': 'application/json'})
    with urllib.request.urlopen(req, timeout=45) as response:
        res = json.load(response)
    if res.get('code') != 0:
        raise RuntimeError('Cloudreve API失败，代码：' + str(res.get('code')) + '，提示：' + str(res.get('msg'))[:200])
    return res.get('data')


def secret():
    for line in (BASE / 'aria2.conf').read_text().splitlines():
        if line.startswith('rpc-secret='):
            return line.split('=', 1)[1]
    raise RuntimeError('RPC凭据缺失')


def rpc(method, args=None, auth=True):
    req = urllib.request.Request('http://127.0.0.1:6800/jsonrpc',
        data=json.dumps({'jsonrpc': '2.0', 'id': 'verify', 'method': 'aria2.' + method,
                         'params': (['token:' + secret()] if auth else []) + (args or [])}).encode(),
        headers={'Content-Type': 'application/json'})
    try:
        with urllib.request.urlopen(req, timeout=15) as response:
            res = json.load(response)
    except urllib.error.HTTPError as error:
        if error.code != 400:
            raise
        res = json.load(error)
    return res


def configure(apply):
    data = api('GET', '/admin/node/1')
    node = data
    assert node['type'] == 'master', '节点身份不符'
    before = copy.deepcopy(node)
    node.setdefault('settings', {})['provider'] = 'aria2'
    node['settings']['aria2'] = {'server': 'http://127.0.0.1:6800', 'token': secret(),
                                'temp_path': str(TEMP), 'options': {}}
    # 保持现有能力、节点状态、用户组和URL安全校验不变。
    print('计划：主节点1补齐Aria2连接及外盘临时目录，不修改其他设置。')
    if apply:
        PRIVATE.mkdir(parents=True, exist_ok=True, mode=0o700)
        backup = PRIVATE / ('node-before-' + time.strftime('%Y%m%d-%H%M%S') + '.json')
        fd = os.open(backup, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
        with os.fdopen(fd, 'w') as out:
            json.dump(before, out)
        version = api('POST', '/admin/node/test/downloader', {'node': node})
        api('PUT', '/admin/node/1', {'node': node})
        after = api('GET', '/admin/node/1')
        assert after['settings']['aria2']['temp_path'] == str(TEMP)
        assert after['capabilities'] == before['capabilities']
        print('节点已保存，通信测试通过，Aria2版本：', version)


def check():
    version = rpc('getVersion')
    assert 'result' in version, 'RPC认证失败'
    rejected = rpc('getVersion', auth=False)
    assert rejected.get('error'), '未认证RPC意外放行'
    node = api('GET', '/admin/node/1')
    assert node['settings']['aria2']['token'] == secret()
    print('通过：认证RPC、未认证拒绝、节点凭据一致。版本：', version['result']['version'])
    print('Cloudreve下载器通信：', api('POST', '/admin/node/test/downloader', {'node': node}))


def main():
    p = argparse.ArgumentParser()
    p.add_argument('action', choices=['configure', 'check'])
    p.add_argument('--apply', action='store_true')
    a = p.parse_args()
    if a.action == 'configure':
        configure(a.apply)
    else:
        check()

if __name__ == '__main__':
    main()