194 lines
10 KiB
Python
194 lines
10 KiB
Python
#!/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')
|
||
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=='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)
|
||
|
||
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()
|