#!/usr/bin/python3
"""Authenticated local task broker. Only the web service UID may connect to its socket."""
import argparse
import getpass
import hashlib
import hmac
import json
import os
from pathlib import Path
try:import pwd
except ImportError:pwd=None
import secrets
import socket
import socketserver
import struct
import re
import datetime as dt
import threading
import time
import task_core as core

ADMIN=Path('/etc/hardware-task-admin.json')
MONITOR_CONFIG=Path('/etc/hardware-watch.json')
SOCKET=Path('/run/hardware-taskd/control.sock')
SESSIONS={};ATTEMPTS={};GUARD=threading.RLock();MUTATE=threading.RLock()
ALLOWED_UID=None

def password_record(password):
    salt=secrets.token_hex(16)
    return {'username':'admin','salt':salt,'hash':hashlib.pbkdf2_hmac('sha256',password.encode(),bytes.fromhex(salt),240000).hex()}

def verify(password):
    try:
        record=json.loads(ADMIN.read_text())
        hashed=hashlib.pbkdf2_hmac('sha256',password.encode(),bytes.fromhex(record['salt']),240000).hex()
        return hmac.compare_digest(record['hash'],hashed)
    except (OSError,ValueError,KeyError):return False

def create_admin():
    if ADMIN.exists():return
    password=secrets.token_urlsafe(18)
    core.atom_write(ADMIN,json.dumps(password_record(password)))
    core.atom_write(Path('/root/hardware-task-admin-login.txt'),
        '定时任务管理员登录信息\n\n网页：http://192.168.1.37:8090/#tasks\n账号：admin\n初始密码：'+password+
        '\n\n此账号可以管理以 root 身份执行的服务器任务，请妥善保存。登录后可修改管理密码。\nSSH root 密码未被修改，也没有作为网页管理密码保存。\n')
    print('初始登录信息已保存到 /root/hardware-task-admin-login.txt（仅 root 可读）。')

def session(token,csrf=None,write=False):
    with GUARD:
        record=SESSIONS.get(token)
        if not record or record['expires']<time.monotonic():
            SESSIONS.pop(token,None);raise core.TaskError('请先登录定时任务管理。',401)
        if write and (not isinstance(csrf,str) or not hmac.compare_digest(record['csrf'],csrf)):
            raise core.TaskError('操作验证已失效，请重新登录后再试。',403)
        return record

def login(payload,client):
    now=time.monotonic()
    with GUARD:
        recent=[x for x in ATTEMPTS.get(client,[]) if now-x<300]
        if len(recent)>=6:raise core.TaskError('登录失败次数较多，请五分钟后重试。',429)
        if payload.get('username')!='admin' or not verify(str(payload.get('password',''))):
            ATTEMPTS[client]=recent+[now]
            core.audit('login_failed',client=client)
            raise core.TaskError('账号或密码不正确。',401)
        ATTEMPTS.pop(client,None)
        if len(SESSIONS)>200:
            for key in list(SESSIONS):
                if SESSIONS[key]['expires']<now:SESSIONS.pop(key,None)
        token=secrets.token_urlsafe(32);csrf=secrets.token_urlsafe(24)
        SESSIONS[token]={'username':'admin','csrf':csrf,'expires':now+8*3600}
    core.audit('login',client=client)
    return {'token':token,'csrf':csrf,'logged_in':True,'username':'admin','max_age':8*3600}

def dispatch(request):
    action=request.get('action','');payload=request.get('payload') or {};token=request.get('token','');client=str(request.get('client','local'))[:100]
    if not isinstance(payload,dict):raise core.TaskError('请求格式不正确。')
    if action=='auth.login':return login(payload,client)
    if action=='auth.status':
        try:record=session(token)
        except core.TaskError:return {'logged_in':False,'configured':ADMIN.exists()}
        return {'logged_in':True,'username':record['username'],'csrf':record['csrf']}
    write=action not in ('tasks.list','tasks.get','tasks.runs','tasks.runlog','tasks.history')
    session(token,request.get('csrf'),write)
    if action=='auth.logout':
        with GUARD:SESSIONS.pop(token,None)
        return {'message':'已退出任务管理。'}
    if action=='auth.password':
        new=str(payload.get('new_password',''))
        if not 12<=len(new)<=200:raise core.TaskError('新密码长度需要 12 至 200 个字符。')
        if not verify(str(payload.get('old_password',''))):raise core.TaskError('原管理密码不正确。',400)
        core.atom_write(ADMIN,json.dumps(password_record(new)))
        with GUARD:SESSIONS.clear()
        core.audit('password_changed',client=client)
        return {'message':'管理密码已修改，请重新登录。'}
    if action=='tasks.list':return core.list_tasks()
    if action=='tasks.get':return core.get_task(payload.get('id'))
    if action=='tasks.runs':return core.runs_for(payload.get('id'))
    if action=='tasks.runlog':return core.read_run(payload.get('run_id'))
    if action=='tasks.history':
        with core.db() as con:rows=con.execute('SELECT * FROM runs ORDER BY started DESC LIMIT 80').fetchall()
        return {'runs':[dict(row) for row in rows]}
    if action=='tasks.preview':
        definition=core.validate_definition(payload)
        calendars=core.calendar_preview(definition)
        next_times=[]
        for output in calendars:
            match=re.search(r'Next elapse:.*?(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2})',output)
            if match:next_times.append(dt.datetime.strptime(match[1],'%Y-%m-%d %H:%M:%S').replace(tzinfo=core.TZ).timestamp())
        if not calendars:
            next_times=[time.time()+int(definition['schedule_lines'][0].split('=')[1][:-1])]
        return {'schedule_label':definition['schedule_label'],'next_run':min(next_times) if next_times else None,'dangerous':definition['dangerous']}
    with MUTATE:
        if action=='monitor.gpu_count':return save_gpu_count(payload,client)
        if action=='tasks.save':return core.save_task(payload,client)
        if action=='tasks.enable':
            if not isinstance(payload.get('enabled'),bool):raise core.TaskError('启用状态格式不正确。')
            return core.set_enabled(payload.get('id'),payload['enabled'],client)
        if action=='tasks.delete':return core.delete_task(payload.get('id'),client)
        if action=='tasks.run':return core.run_now(payload.get('id'),client)
        if action=='tasks.stop':return core.stop_task(payload.get('id'),client)
    raise core.TaskError('不支持这个任务操作。',404)

def save_gpu_count(payload,client=''):
    count=payload.get('expected_gpu_count')
    if type(count) is not int or count not in (1,2):
        raise core.TaskError('请选择使用 1 张或 2 张显卡。')
    config=json.loads(MONITOR_CONFIG.read_text())
    if not isinstance(config,dict):raise core.TaskError('硬件监控配置格式不正确。',500)
    config.update(expected_gpu_count=count,expected_gpu_uuids=[])
    mode=MONITOR_CONFIG.stat().st_mode & 0o777
    core.atom_write(MONITOR_CONFIG,json.dumps(config,ensure_ascii=False,indent=2)+'\n',mode)
    # Replacement drops any ACL on the old file. Grant the dashboard's service
    # group read access while keeping configuration writes restricted to root.
    if pwd is not None and os.geteuid()==0:
        import grp
        os.chown(MONITOR_CONFIG,0,grp.getgrnam('hardware-watch').gr_gid)
        os.chmod(MONITOR_CONFIG,0o640)
    core.audit('gpu_count_changed',client=client,detail=f'expected_gpu_count={count}; identity=count')
    refresh=core.execute(['systemctl','start','--no-block','hardware-watch.service'],check=False)
    return {'expected_gpu_count':count,'message':f'已设置使用 {count} 张显卡；按数量检查，允许更换显卡。',
            'refresh_pending':refresh.returncode==0}

class Handler(socketserver.StreamRequestHandler):
    def handle(self):
        self.request.settimeout(25)
        try:
            _,uid,_=struct.unpack('3i',self.request.getsockopt(socket.SOL_SOCKET,socket.SO_PEERCRED,12))
            if uid not in (0,ALLOWED_UID):raise core.TaskError('不允许连接任务服务。',403)
            raw=self.rfile.readline(131073)
            if len(raw)>131072:raise core.TaskError('请求内容过大。',413)
            request=json.loads(raw)
            response={'ok':True,'data':dispatch(request)}
        except core.TaskError as exc:response={'ok':False,'error':str(exc),'status':exc.status}
        except Exception as exc:
            print('task broker error:',type(exc).__name__,str(exc),flush=True)
            response={'ok':False,'error':'任务操作未完成，请查看任务管理服务日志。','status':500}
        try:self.wfile.write(json.dumps(response,ensure_ascii=False).encode()+b'\n')
        except (BrokenPipeError,ConnectionResetError):pass

class Server(getattr(socketserver,'ThreadingUnixStreamServer',socketserver.ThreadingTCPServer)):
    daemon_threads=True

def maintenance():
    while True:
        try:
            with core.db() as con:
                rows=con.execute("SELECT * FROM runs WHERE state='running'").fetchall()
            boot=Path('/proc/sys/kernel/random/boot_id').read_text().strip()
            for row in rows:
                if row['boot']!=boot or (time.time()-row['started']>30 and not core.running(row['task_id'])):
                    with core.db() as con:
                        con.execute("UPDATE runs SET state='interrupted',finished=? WHERE id=? AND state='running'",(time.time(),row['id']))
                    with (core.LOGS/(row['id']+'.log')).open('ab') as f:
                        f.write(('\n[执行被系统重启或进程中断，没有取得完整退出结果。]\n').encode())
            cutoff=time.time()-90*86400
            with core.db() as con:
                old=con.execute('SELECT id FROM runs WHERE started<? AND state!=?',(cutoff,'running')).fetchall()
                for row in old:
                    (core.LOGS/(row['id']+'.log')).unlink(missing_ok=True)
                    for suffix in ('.sh','.py'):(core.BASE/'scripts'/(row['id']+suffix)).unlink(missing_ok=True)
                con.execute("DELETE FROM runs WHERE started<? AND state!='running'",(cutoff,))
                con.execute('DELETE FROM audit WHERE at<?',(cutoff,))
                protected={x['id'] for x in con.execute("SELECT id FROM runs WHERE state='running'")}
                files=[p for folder in (core.LOGS,core.BASE/'scripts') for p in folder.iterdir()
                       if p.is_file() and re.fullmatch(r'[0-9a-f]{32}\.(log|sh|py)',p.name)]
                files.sort(key=lambda p:p.stat().st_mtime)
                total=sum(p.stat().st_size for p in files)
                for path in files:
                    if total<=256*1024*1024:break
                    if path.stem in protected:continue
                    total-=path.stat().st_size;path.unlink()
                    if path.suffix=='.log':con.execute('UPDATE runs SET log_purged=1 WHERE id=?',(path.stem,))
        except Exception as exc:print('task maintenance:',type(exc).__name__,str(exc),flush=True)
        time.sleep(60)

def main():
    global ALLOWED_UID
    parser=argparse.ArgumentParser();parser.add_argument('mode',choices=['serve','init-admin','password'],default='serve',nargs='?');args=parser.parse_args()
    core.init_store()
    if args.mode=='init-admin':create_admin();return
    if args.mode=='password':
        password=getpass.getpass('新管理密码（至少12字符）：')
        if len(password)<12 or password!=getpass.getpass('再次输入：'):raise SystemExit('密码太短或两次输入不一致。')
        core.atom_write(ADMIN,json.dumps(password_record(password)));print('密码已重置；请重启 hardware-taskd 服务使旧登录失效。');return
    if not ADMIN.exists():raise SystemExit('请先运行 taskd.py init-admin。')
    ALLOWED_UID=pwd.getpwnam('hardware-dashboard').pw_uid
    SOCKET.parent.mkdir(parents=True,exist_ok=True)
    SOCKET.unlink(missing_ok=True)
    server=Server(str(SOCKET),Handler);os.chmod(SOCKET,0o660)
    threading.Thread(target=maintenance,daemon=True).start()
    print('Task broker ready on local Unix socket.',flush=True)
    server.serve_forever()

if __name__=='__main__':main()