Files
ai-worker-platform/enterprise.py
T

339 lines
13 KiB
Python
Raw Normal View History

# -*- 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