hardware-dashboard/task_exec.py

#!/usr/bin/python3
"""Run one approved, root-owned task definition; record bounded output and result."""
import fcntl
import hashlib
import json
import os
import selectors
import shlex
import signal
import subprocess
import sys
import time
import uuid
import task_core as core

LIMIT=512*1024

def main(task_id):
    task=core.get_task(task_id)
    definition=task['definition']
    lock=(core.BASE/(task_id+'.lock')).open('a')
    try:fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
    except BlockingIOError:return 0
    run_id=uuid.uuid4().hex;started=time.time();start_mono=time.monotonic()
    boot=core.read_boot() if hasattr(core,'read_boot') else open('/proc/sys/kernel/random/boot_id').read().strip()
    with core.db() as con:
        requested=con.execute('SELECT * FROM requests WHERE task_id=?',(task_id,)).fetchone()
        trigger='manual' if requested and requested['expires']>=started else 'scheduled'
        con.execute('DELETE FROM requests WHERE task_id=?',(task_id,))
        body=definition['body']
        fingerprint=hashlib.sha256((body if definition['mode']!='file' else definition['path']+' '+definition['arguments']).encode()).hexdigest()
        con.execute('INSERT INTO runs(id,task_id,task_name,started,state,trigger,boot,script_hash,revision) VALUES(?,?,?,?,?,?,?,?,?)',
                    (run_id,task_id,task['name'],started,'running',trigger,boot,fingerprint,task['revision']))
    cancelled=[False]
    def stop(*_):cancelled[0]=True
    signal.signal(signal.SIGTERM,stop);signal.signal(signal.SIGINT,stop)
    state='failed';exit_code=None;proc=None;written=0;truncated=False
    log_path=core.LOGS/(run_id+'.log')
    with log_path.open('wb',buffering=0) as output:
        os.chmod(log_path,0o600)
        def write(data):
            nonlocal written,truncated
            if isinstance(data,str):data=data.encode('utf-8')
            available=max(0,LIMIT-written)
            if available:
                piece=data[:available];output.write(piece);written+=len(piece)
            if len(data)>available:truncated=True
        write(f'开始：{core.iso(started)}\n任务：{task["name"]}\n触发：{"手动" if trigger=="manual" else "定时"}\n执行身份：root\n修订：{task["revision"]}\n'+'─'*32+'\n')
        try:
            if definition['mode']=='command':
                argv=['/bin/bash','-e','-o','pipefail','-c',body]
            elif definition['mode']=='script':
                script=core.BASE/'scripts'/(run_id+('.py' if definition['interpreter']=='python3' else '.sh'))
                core.atom_write(script,body+'\n')
                argv=['/usr/bin/python3',str(script)] if definition['interpreter']=='python3' else ['/bin/bash','-e','-o','pipefail',str(script)]
            else:
                argv=[definition['path'],*shlex.split(definition['arguments'])]
                if definition['interpreter']=='bash':argv=['/bin/bash','-e','-o','pipefail',*argv]
                elif definition['interpreter']=='python3':argv=['/usr/bin/python3',*argv]
            env={'HOME':'/root','USER':'root','LOGNAME':'root','LANG':'C.UTF-8','TZ':'Asia/Shanghai',
                 'PATH':'/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/snap/bin',
                 'HARDWARE_TASK_ID':task_id,'HARDWARE_TASK_RUN_ID':run_id}
            proc=subprocess.Popen(argv,cwd=definition['cwd'],env=env,stdout=subprocess.PIPE,stderr=subprocess.STDOUT,start_new_session=True)
            selector=selectors.DefaultSelector();selector.register(proc.stdout,selectors.EVENT_READ)
            deadline=start_mono+definition['timeout'];exited_at=None
            while selector.get_map() or proc.poll() is None:
                if cancelled[0] or time.monotonic()>=deadline:
                    state='cancelled' if cancelled[0] else 'timeout'
                    try:os.killpg(proc.pid,signal.SIGTERM)
                    except ProcessLookupError:pass
                    try:proc.wait(timeout=3)
                    except subprocess.TimeoutExpired:
                        try:os.killpg(proc.pid,signal.SIGKILL)
                        except ProcessLookupError:pass
                        proc.wait(timeout=3)
                    break
                for key,_ in selector.select(.25):
                    chunk=os.read(key.fileobj.fileno(),32768)
                    if chunk:write(chunk)
                    else:selector.unregister(key.fileobj)
                if proc.poll() is not None:
                    exited_at=exited_at or time.monotonic()
                    if time.monotonic()-exited_at>1:break
            selector.close();proc.stdout.close()
            exit_code=proc.wait(timeout=3)
            if state not in ('cancelled','timeout'):state='success' if exit_code==0 else 'failed'
        except Exception as exc:
            write('\n执行错误：'+str(exc)+'\n');state='failed'
            if proc and proc.poll() is None:
                try:os.killpg(proc.pid,signal.SIGKILL);proc.wait(timeout=3)
                except (ProcessLookupError,subprocess.TimeoutExpired):pass
        elapsed=time.monotonic()-start_mono
        tail=f'\n'+('─'*32)+f'\n结束：{core.iso()}\n结果：{state}\n退出码：{exit_code}\n耗时：{elapsed:.2f} 秒\n'
        if truncated:output.write('\n[输出超过 512 KiB，后续内容已省略。]\n'.encode())
        output.write(tail.encode())
        os.fsync(output.fileno())
    with core.db() as con:
        con.execute('UPDATE runs SET finished=?,state=?,exit_code=?,log_bytes=?,truncated=?,duration=? WHERE id=?',
                    (time.time(),state,exit_code,log_path.stat().st_size,int(truncated),elapsed,run_id))
    lock.close()
    return 0 if state=='success' else 1

if __name__=='__main__':
    raise SystemExit(main(core.validate_id(sys.argv[1])))