Files
ai-worker-platform/enterprise.py
T
hz4th_coder 9a95b71631 V2 开发:多 Agent 协作(主管/评审/辩论) + 自动评估体系 + 模板市场 + 企业版(SSO/RBAC/审计/合规)
- agents.py: supervisor/review/debate 三模式编排引擎,步骤流水全记录
- eval.py: eval 数据集 + LLM-as-Judge 评分 + 任务产出自动沉淀 + Worker 排行榜
- templates.py: 任务/项目/团队三类模板,占位符渲染(JSON安全转义),一键应用
- enterprise.py: 用户体系/RBAC、OIDC+LDAP SSO、审计日志、数据导出/保留期/PII脱敏
- db.py: V2 表结构 + 迁移 + 启动恢复(stale running→failed)
- app.py: V2 API 路由 + 登录升级(账号体系)+ 审计钩子
- 前端: 新增多Agent协作/自动评估/模板市场/企业版 4 页面
- Dockerfile: 企业版私有化部署
2026-08-12 12:21:17 +08:00

340 lines
13 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# -*- coding: utf-8 -*-
"""
V2 企业版能力
- 用户体系 + RBACadmin / member / auditor
- SSOOIDC 授权码模式 + LDAP 绑定(可选依赖),本地口令兜底
- 审计日志:全 API 关键操作留痕(actor/action/target/ip
- 合规:数据导出(全量 JSON / 审计 CSV)、保留期清理、PII 掩码
"""
import csv
import io
import json
import time
import uuid
import db
# ---------------------------------------------------------------------------
# 审计
# ---------------------------------------------------------------------------
PII_PATTERNS = [
(r'[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}', '<email>'),
(r'\b1[3-9]\d{9}\b', '<phone>'),
(r'\b\d{17}[\dXx]\b', '<idcard>'),
]
def mask_pii(text):
"""合规:敏感信息掩码"""
if not text:
return text
import re
for pat, rep in PII_PATTERNS:
text = re.sub(pat, rep, text)
return text
def audit(actor, action, target='', detail='', ip='', user_agent=''):
"""写审计日志(失败不影响主流程)"""
try:
if db.get_setting('compliance_mask_pii', '0') == '1':
detail = mask_pii(detail)
db.w(
'INSERT INTO audit_logs (actor, action, target, detail, ip, user_agent, created_at) '
'VALUES (?,?,?,?,?,?,?)',
(str(actor)[:100], str(action)[:100], str(target)[:200], str(detail)[:2000],
str(ip)[:64], str(user_agent)[:200], db.now()))
except Exception:
pass
def current_actor():
"""从请求上下文推断操作者(由 app 注入 request-local 变量)"""
import flask
try:
req = flask.request
actor = getattr(flask.g, 'auth_actor', None)
if actor:
return actor
hdr = req.headers.get('Authorization', '')
if hdr.startswith('Bearer '):
r = db.q('SELECT name FROM api_tokens WHERE token=?', (hdr[7:].strip(),), one=True)
return f'token:{r["name"]}' if r else 'token:?'
username = flask.session.get('username')
if username:
return username
return 'anonymous'
except Exception:
return 'anonymous'
def audit_auto(action, target='', detail='', save_body_keys=None):
"""装饰器版自动审计:包装 flask 路由"""
import flask
import functools
def deco(fn):
@functools.wraps(fn)
def wrapper(*args, **kwargs):
resp = fn(*args, **kwargs)
try:
detail_text = ''
if save_body_keys and flask.request.method in ('POST', 'PUT'):
try:
body = flask.request.get_json(silent=True) or {}
detail_text = ' '.join(f'{k}={body.get(k)}' for k in save_body_keys if k in body)
except Exception:
pass
audit(current_actor(), action, target or (flask.request.path or ''),
detail_text, flask.request.remote_addr or '',
flask.request.headers.get('User-Agent', ''))
except Exception:
pass
return resp
return wrapper
return deco
# ---------------------------------------------------------------------------
# 用户 / RBAC
# ---------------------------------------------------------------------------
def get_user(username):
return db.q('SELECT * FROM users WHERE username=?', (username,), one=True)
def role_of(username):
u = get_user(username)
return u['role'] if u else 'anonymous'
def is_admin(username):
return role_of(username) == 'admin'
def create_local_user(username, password, display_name='', role='member'):
if get_user(username):
return None, '用户已存在'
uid = db.w(
'INSERT INTO users (username, password_hash, display_name, role, source, status, created_at) '
'VALUES (?,?,?,?,?,?,?)',
(username, db.hash_password(password), display_name or username, role, 'local', 'active', db.now()))
return uid, None
def verify_local(username, password):
u = get_user(username)
if not u or u['source'] != 'local' or u['status'] != 'active':
return None
if db.verify_password(password, u['password_hash']):
db.w('UPDATE users SET last_login_at=? WHERE id=?', (db.now(), u['id']))
return u
return None
# ---------------------------------------------------------------------------
# SSOOIDC(授权码)+ LDAP(可选)
# ---------------------------------------------------------------------------
def sso_config():
cfg = {}
for k in ('oidc_enabled', 'oidc_name', 'oidc_discovery_url', 'oidc_client_id',
'oidc_client_secret', 'oidc_redirect_uri', 'oidc_scope', 'oidc_admin_group',
'ldap_enabled', 'ldap_url', 'ldap_base_dn', 'ldap_bind_dn', 'ldap_bind_password',
'ldap_user_filter', 'sso_auto_provision'):
cfg[k] = db.get_setting(k, '')
return cfg
def save_sso_config(data):
keys = list(data.keys())
for k in keys:
if k in ('oidc_client_secret', 'ldap_bind_password') and not data[k]:
continue # 留空不覆盖已保存的密钥
db.set_setting(k, str(data[k]))
def oidc_discovery():
"""读取 OIDC discovery 文档,返回端点字典"""
import requests
url = sso_config().get('oidc_discovery_url', '').strip().rstrip('/')
if not url:
raise ValueError('未配置 OIDC discovery URL')
r = requests.get(url, timeout=15)
if r.status_code != 200:
raise ValueError(f'Discovery 请求失败({r.status_code})')
return r.json()
def oidc_authorize_url(state):
"""生成授权跳转 URL"""
import urllib.parse
cfg = sso_config()
disc = oidc_discovery()
params = {
'response_type': 'code',
'client_id': cfg['oidc_client_id'],
'redirect_uri': cfg['oidc_redirect_uri'],
'scope': cfg.get('oidc_scope') or 'openid profile email',
'state': state,
}
return disc.get('authorization_endpoint') + '?' + urllib.parse.urlencode(params)
def oidc_exchange(code):
"""用授权码换 token + 用户信息"""
import requests
cfg = sso_config()
disc = oidc_discovery()
tok = requests.post(disc.get('token_endpoint'), data={
'grant_type': 'authorization_code',
'code': code,
'redirect_uri': cfg['oidc_redirect_uri'],
'client_id': cfg['oidc_client_id'],
'client_secret': cfg['oidc_client_secret'],
}, timeout=15)
if tok.status_code != 200:
raise ValueError(f'Token 交换失败({tok.status_code}): {tok.text[:200]}')
token_data = tok.json()
id_token = token_data.get('id_token', '')
userinfo = {}
# 优先 userinfo 端点
if token_data.get('access_token'):
ui = requests.get(disc.get('userinfo_endpoint'), headers={
'Authorization': f"Bearer {token_data['access_token']}"}, timeout=15)
if ui.status_code == 200:
userinfo = ui.json()
# 解析 id_token payload 兜底
if not userinfo and id_token:
import base64
try:
payload = id_token.split('.')[1]
payload += '=' * (-len(payload) % 4)
userinfo = json.loads(base64.urlsafe_b64decode(payload))
except Exception:
pass
return userinfo
def sso_login(userinfo):
"""SSO 登录回调:查找或自动开通用户"""
cfg = sso_config()
username = userinfo.get('preferred_username') or userinfo.get('email') or userinfo.get('sub') or ''
email = userinfo.get('email', '')
display = userinfo.get('name') or userinfo.get('display_name') or username
groups = userinfo.get('groups') or userinfo.get('roles') or []
if not username:
return None, '无法从 SSO 响应中解析用户名'
u = get_user(username)
if not u:
if cfg.get('sso_auto_provision') != '1':
return None, '用户未开通(自动开通未启用),请联系管理员'
role = 'admin' if cfg.get('oidc_admin_group') and cfg['oidc_admin_group'] in groups else 'member'
db.w(
'INSERT INTO users (username, password_hash, display_name, role, source, status, created_at) '
'VALUES (?,?,?,?,?,?,?)',
(username, '', display, role, 'oidc', 'active', db.now()))
u = get_user(username)
elif u['status'] != 'active':
return None, '账号已停用'
elif u['source'] != 'oidc':
return None, f'用户名 {username} 已被本地账号占用'
db.w('UPDATE users SET last_login_at=? WHERE id=?', (db.now(), u['id']))
return u, None
def ldap_authenticate(username, password):
"""LDAP 绑定认证(依赖 ldap3,未安装时返回 None"""
try:
from ldap3 import Server, Connection, ALL
except ImportError:
return None, '未安装 ldap3,无法使用 LDAP SSO'
cfg = sso_config()
try:
server = Server(cfg['ldap_url'], get_info=ALL)
conn = Connection(server, user=cfg['ldap_bind_dn'], password=cfg['ldap_bind_password'],
auto_bind=True)
user_filter = cfg.get('ldap_user_filter') or '(uid={username})'
conn.search(cfg['ldap_base_dn'], user_filter.format(username=username), attributes=['cn', 'mail', 'displayName'])
if not conn.entries:
return None, 'LDAP 中未找到该用户'
entry = conn.entries[0]
user_dn = entry.entry_dn
conn.unbind()
conn2 = Connection(server, user=user_dn, password=password, auto_bind=True)
conn2.unbind()
return {'username': username,
'display': str(entry.displayName.value) if hasattr(entry, 'displayName') else username,
'email': str(entry.mail.value) if hasattr(entry, 'mail') else ''}, None
except Exception as e:
return None, f'LDAP 认证失败: {e}'
def ldap_login(username, password):
info, err = ldap_authenticate(username, password)
if err:
return None, err
u = get_user(username)
if not u:
if sso_config().get('sso_auto_provision') != '1':
return None, '用户未开通,请联系管理员'
db.w(
'INSERT INTO users (username, password_hash, display_name, role, source, status, created_at) '
'VALUES (?,?,?,?,?,?,?)',
(username, '', info.get('display') or username, 'member', 'ldap', 'active', db.now()))
u = get_user(username)
elif u['status'] != 'active':
return None, '账号已停用'
db.w('UPDATE users SET last_login_at=? WHERE id=?', (db.now(), u['id']))
return u, None
# ---------------------------------------------------------------------------
# 合规:导出 / 保留期
# ---------------------------------------------------------------------------
def export_all():
"""全量数据导出(JSON"""
tables = ['projects', 'workers', 'tasks', 'task_logs', 'cost_records', 'documents',
'agent_runs', 'agent_steps', 'eval_datasets', 'eval_cases', 'eval_runs',
'eval_results', 'templates', 'users', 'audit_logs']
out = {'exported_at': time.strftime('%Y-%m-%d %H:%M:%S'),
'platform': 'ai-worker-platform', 'version': 'v2.0.0'}
for t in tables:
try:
out[t] = db.q(f'SELECT * FROM {t}')
except Exception:
out[t] = []
return out
def export_audit_csv():
"""审计日志导出 CSV"""
rows = db.q('SELECT * FROM audit_logs ORDER BY id DESC LIMIT 10000')
buf = io.StringIO()
w = csv.writer(buf)
w.writerow(['ID', '时间', '操作者', '动作', '目标', '详情', 'IP', 'UA'])
for r in rows:
w.writerow([r['id'], time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(r['created_at'])),
r['actor'], r['action'], r['target'], r['detail'], r['ip'], r['user_agent']])
return buf.getvalue()
def apply_retention():
"""合规:按保留期清理审计日志 / 协作运行 / 评估结果(每天可跑一次)"""
days = int(db.get_setting('compliance_retention_days', '0') or 0)
if days <= 0:
return {'cleaned': 0, 'note': '未配置保留期(0=永久保留)'}
cutoff = db.now() - days * 86400
cleaned = 0
for table in ('audit_logs', 'agent_steps', 'agent_runs', 'eval_results', 'eval_runs'):
try:
cur = db.w(f'DELETE FROM {table} WHERE created_at<?', (cutoff,))
cleaned += cur or 0
except Exception:
pass
return {'cleaned': cleaned, 'retention_days': days}
def generate_consent_token():
"""生成数据使用同意记录 token(审计用途)"""
tok = uuid.uuid4().hex[:12]
audit('system', 'compliance.consent', '数据使用同意', f'consent_token={tok}')
return tok