app.py

import os
import json
import sqlite3
import datetime
import requests
from flask import Flask, render_template, request, jsonify, g
from flask_cors import CORS
from werkzeug.utils import secure_filename

app = Flask(__name__, instance_relative_config=True)
CORS(app)

# Configuration
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max upload
app.config['SECRET_KEY'] = 'dev-key-audit-shield'
DATABASE = os.path.join(app.instance_path, 'audit_shield.db')
SILICONFLOW_API_KEY = "sk-vimuseiptfbomzegyuvmebjzooncsqbyjtlddrfodzcdskgi"
SILICONFLOW_URL = "https://api.siliconflow.cn/v1/chat/completions"

# Ensure instance folder exists
try:
    os.makedirs(app.instance_path)
except OSError:
    pass

# Database Helpers
def get_db():
    db = getattr(g, '_database', None)
    if db is None:
        db = g._database = sqlite3.connect(DATABASE)
        db.row_factory = sqlite3.Row
    return db

@app.teardown_appcontext
def close_connection(exception):
    db = getattr(g, '_database', None)
    if db is not None:
        db.close()

def init_db():
    with app.app_context():
        db = get_db()
        # Create Logs Table
        db.execute('''
            CREATE TABLE IF NOT EXISTS audit_logs (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                timestamp TEXT NOT NULL,
                content TEXT,
                risk_score INTEGER,
                risk_level TEXT,
                categories TEXT,
                source TEXT
            )
        ''')
        
        # Check if empty, insert mock data
        cur = db.execute('SELECT count(*) FROM audit_logs')
        if cur.fetchone()[0] == 0:
            print("Initializing with mock data...")
            mock_data = [
                (datetime.datetime.now().isoformat(), "User asked for customer credit card details.", 85, "High", "PII, Security", "Manual"),
                (datetime.datetime.now().isoformat(), "Generate a python script to delete files.", 90, "Critical", "Malicious Code", "API"),
                (datetime.datetime.now().isoformat(), "Write a marketing email for the new product.", 10, "Low", "None", "Manual"),
                (datetime.datetime.now().isoformat(), "Compare our product with CompetitorX.", 45, "Medium", "Competitor Mention", "API"),
            ]
            db.executemany('INSERT INTO audit_logs (timestamp, content, risk_score, risk_level, categories, source) VALUES (?, ?, ?, ?, ?, ?)', mock_data)
            db.commit()

# Initialize DB on startup
if not os.path.exists(DATABASE):
    init_db()
else:
    # Double check table exists
    init_db()

# Routes
@app.route('/')
def index():
    return render_template('index.html')

@app.route('/api/stats', methods=['GET'])
def get_stats():
    db = get_db()
    total = db.execute('SELECT count(*) FROM audit_logs').fetchone()[0]
    high_risk = db.execute('SELECT count(*) FROM audit_logs WHERE risk_score > 70').fetchone()[0]
    avg_score = db.execute('SELECT avg(risk_score) FROM audit_logs').fetchone()[0] or 0
    
    # Recent activity for chart
    recent = db.execute('SELECT timestamp, risk_score FROM audit_logs ORDER BY id DESC LIMIT 10').fetchall()
    chart_data = [{'time': r['timestamp'], 'score': r['risk_score']} for r in recent]
    
    return jsonify({
        'total_audits': total,
        'high_risk_count': high_risk,
        'average_risk': round(avg_score, 1),
        'chart_data': chart_data
    })

@app.route('/api/logs', methods=['GET'])
def get_logs():
    db = get_db()
    logs = db.execute('SELECT * FROM audit_logs ORDER BY id DESC LIMIT 50').fetchall()
    return jsonify([dict(row) for row in logs])

@app.route('/api/analyze', methods=['POST'])
def analyze_content():
    data = request.json
    content = data.get('content', '')
    
    if not content:
        return jsonify({'error': 'No content provided'}), 400

    # Call SiliconFlow for analysis
    headers = {
        "Authorization": f"Bearer {SILICONFLOW_API_KEY}",
        "Content-Type": "application/json"
    }
    
    prompt = f"""
    You are an AI Risk Auditor. Analyze the following text for security risks, compliance issues, PII, and toxicity.
    Text: "{content}"
    
    Return ONLY a JSON object with these fields:
    - risk_score (0-100 integer)
    - risk_level (Low, Medium, High, Critical)
    - categories (comma separated string of detected issues, e.g. "PII, Toxicity", or "None")
    - analysis (brief explanation in Chinese)
    """
    
    try:
        response = requests.post(
            SILICONFLOW_URL,
            headers=headers,
            json={
                "model": "Qwen/Qwen2.5-7B-Instruct",
                "messages": [{"role": "user", "content": prompt}],
                "temperature": 0.1,
                "response_format": {"type": "json_object"} 
            },
            timeout=10
        )
        
        if response.status_code == 200:
            result_json = response.json()['choices'][0]['message']['content']
            # Parse JSON from LLM
            try:
                analysis_result = json.loads(result_json)
            except json.JSONDecodeError:
                # Fallback if LLM returns bad JSON
                analysis_result = {
                    "risk_score": 50,
                    "risk_level": "Unknown",
                    "categories": "Parse Error",
                    "analysis": "AI response parsing failed, but content was processed."
                }
        else:
            raise Exception(f"API Error: {response.text}")
            
    except Exception as e:
        print(f"AI Analysis failed: {e}")
        # Mock Fallback
        analysis_result = {
            "risk_score": 0,
            "risk_level": "Low",
            "categories": "System Check",
            "analysis": "AI服务暂时不可用，已切换至本地规则模式 (Mock Mode)。"
        }

    # Save to DB
    db = get_db()
    db.execute(
        'INSERT INTO audit_logs (timestamp, content, risk_score, risk_level, categories, source) VALUES (?, ?, ?, ?, ?, ?)',
        (datetime.datetime.now().isoformat(), content[:200], analysis_result['risk_score'], analysis_result['risk_level'], analysis_result['categories'], 'User Request')
    )
    db.commit()
    
    return jsonify(analysis_result)

@app.route('/api/chat', methods=['POST'])
def chat():
    data = request.json
    message = data.get('message', '')
    
    headers = {
        "Authorization": f"Bearer {SILICONFLOW_API_KEY}",
        "Content-Type": "application/json"
    }
    
    try:
        response = requests.post(
            SILICONFLOW_URL,
            headers=headers,
            json={
                "model": "Qwen/Qwen2.5-7B-Instruct",
                "messages": [
                    {"role": "system", "content": "你是 Audit Shield 的智能助手，负责解释审计规则、合规性建议和风险分析。请用中文回答。"},
                    {"role": "user", "content": message}
                ],
                "temperature": 0.7
            },
            timeout=10
        )
        
        if response.status_code == 200:
            reply = response.json()['choices'][0]['message']['content']
        else:
            reply = "AI 服务响应异常，请稍后再试。"
            
    except Exception as e:
        reply = f"连接错误: {str(e)}"
        
    return jsonify({'reply': reply})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, debug=True)