app.py

"""
IdentityAI Web Application
Modern Flask-based web interface for threat detection
Features: Dark/Light themes, real-time updates, dynamic visualizations,
          event explorer, user detail drill-down, activity log
"""

from flask import Flask, render_template, jsonify, request, send_file, Response  # pyre-ignore
from flask_cors import CORS  # pyre-ignore
from core.database import init_db, clear_db, DB_NAME  # pyre-ignore
from core.parser import parse_linux_auth, parse_journalctl_json, parse_windows_event, parse_macos_log  # pyre-ignore
from core.normalizer import store_event  # pyre-ignore
from core.detector import run_detection, get_user_details  # pyre-ignore
import os
import sys
import sqlite3
import json
import csv
import io
import subprocess
from datetime import datetime
import threading

app = Flask(__name__)
CORS(app)

# Global status tracking
status = {
    'message': 'Ready',
    'state': 'idle',  # idle, processing, error
    'progress': 0,
    'last_update': datetime.now().isoformat()
}

# In-memory activity log (last 200 entries)
activity_log = []
MAX_LOG_ENTRIES = 200

def add_log(message, level='info'):
    """Add a message to the activity log"""
    entry = {
        'timestamp': datetime.now().strftime('%H:%M:%S'),
        'message': message,
        'level': level
    }
    activity_log.append(entry)
    if len(activity_log) > MAX_LOG_ENTRIES:
        activity_log.pop(0)

# Initialize database on startup
init_db()
add_log('Database initialized', 'success')
add_log('IdentityAI web server started', 'info')

@app.route('/')
def index():
    """Serve the main UI"""
    return render_template('index.html')

@app.route('/api/status', methods=['GET'])
def get_status():
    """Get current system status"""
    try:
        conn = sqlite3.connect(DB_NAME)
        cursor = conn.cursor()

        # Get database statistics
        total_events = cursor.execute("SELECT COUNT(*) FROM events").fetchone()[0]
        total_users = cursor.execute("SELECT COUNT(DISTINCT username) FROM events").fetchone()[0]
        failed_events = cursor.execute("SELECT COUNT(*) FROM events WHERE success = 0").fetchone()[0]

        conn.close()

        success_rate = round((total_events - failed_events) / total_events * 100, 1) if total_events > 0 else 0

        return jsonify({
            'status': status,
            'stats': {
                'total_events': total_events,
                'total_users': total_users,
                'failed_events': failed_events,
                'success_rate': success_rate
            }
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/ingest', methods=['POST'])
def ingest_logs():
    """Ingest logs from specified directory"""
    global status

    data = request.get_json()
    log_dir = data.get('directory', 'sample_logs')

    def ingest_task():
        global status
        try:
            status['state'] = 'processing'
            status['message'] = 'Ingesting logs...'
            status['progress'] = 0
            add_log(f'Starting log ingestion from: {log_dir}', 'info')

            if not os.path.exists(log_dir) or not os.path.isdir(log_dir):
                status['state'] = 'error'
                status['message'] = f'Directory not found: {log_dir}'
                add_log(f'Directory not found: {log_dir}', 'error')
                return

            files = [f for f in os.listdir(log_dir) if os.path.isfile(os.path.join(log_dir, f))]

            if not files:
                status['state'] = 'error'
                status['message'] = 'No log files found'
                add_log('No log files found in directory', 'error')
                return

            add_log(f'Found {len(files)} log file(s)', 'info')
            total_events = 0

            for i, file in enumerate(files):
                filepath = os.path.join(log_dir, file)
                file_events = 0

                try:
                    with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
                        for line in f:
                            event = parse_linux_auth(line)
                            if event:
                                store_event(event)
                                file_events += 1  # pyre-ignore
                                total_events += 1  # pyre-ignore
                except Exception as e:
                    add_log(f'Error processing {file}: {e}', 'error')
                    continue

                status['progress'] = int((i + 1) / len(files) * 100)
                status['message'] = f'Processing {file}... ({file_events} events)'
                if file_events > 0:
                    add_log(f'✓ {file}: {file_events} events', 'success')

            status['state'] = 'idle'
            status['message'] = f'Ingested {total_events} events from {len(files)} files'
            status['progress'] = 100
            add_log(f'Ingestion complete: {total_events} events from {len(files)} files', 'success')

        except Exception as e:
            status['state'] = 'error'
            status['message'] = f'Ingestion failed: {str(e)}'
            add_log(f'Ingestion failed: {str(e)}', 'error')
        finally:
            status['last_update'] = datetime.now().isoformat()

    # Run in background thread
    thread = threading.Thread(target=ingest_task, daemon=True)
    thread.start()

    return jsonify({'message': 'Ingestion started'})

@app.route('/api/scan', methods=['POST'])
def scan_system():
    """Scan active host system logs via journalctl"""
    global status

    def scan_task():
        global status
        total_events = 0
        valid_events = 0
        try:
            status['state'] = 'processing'
            status['message'] = 'Scanning live system logs...'
            status['progress'] = 0
            # Build OS-specific command
            if sys.platform == 'win32':
                cmd = ['powershell', '-NoProfile', '-Command', 
                       "Get-WinEvent -FilterHashtable @{LogName='Security'; ID=4624,4625} -MaxEvents 1000 -ErrorAction SilentlyContinue | Select-Object TimeCreated, Id, Message | ConvertTo-Json -Compress"]
                add_log('Connecting to Windows Event Log...', 'info')
            elif sys.platform == 'darwin':
                cmd = ['log', 'show', '--predicate', 'process == "sshd" or process == "login" or process == "sudo" or process == "su"', '--style', 'json', '--last', '7d']
                add_log('Connecting to macOS unified logs...', 'info')
            else:
                cmd = [
                    'journalctl', '-o', 'json',
                    '_COMM=sudo', '_COMM=su', '_COMM=sshd', '_COMM=login', '_COMM=polkitd',
                    '--since', '7 days ago'
                ]
                add_log('Connecting to system journal...', 'info')
            
            add_log('Running log extraction...', 'info')
            
            process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            out, err = process.communicate()
            
            # Process output 
            if sys.platform in ('win32', 'darwin'):
                try:
                    events_data = json.loads(out) if out.strip() else []
                    if isinstance(events_data, dict):
                        events_data = [events_data]
                        
                    for data in events_data:
                        total_events += 1  # pyre-ignore
                        event = parse_windows_event(data) if sys.platform == 'win32' else parse_macos_log(data)
                        if event:
                            store_event(event)
                            valid_events += 1  # pyre-ignore
                except json.JSONDecodeError:
                    add_log(f'Failed to parse JSON output from OS tool', 'error')
            else:
                # Process Linux journalctl line by line processing
                for line in out.splitlines():
                    if not line.strip(): continue
                    total_events += 1  # pyre-ignore
                    event = parse_journalctl_json(line)
                    if event:
                        store_event(event)
                        valid_events += 1  # pyre-ignore

            if process.returncode != 0 and sys.platform != 'win32':
                # Windows might return non-zero if no events found
                err_str = err if err else ""
                add_log(f'OS Tool returned error code {process.returncode}: {err_str[:200]}', 'warning')  # pyre-ignore

            status['state'] = 'idle'
            status['message'] = f'Scan complete: Ingested {valid_events} authentication events'
            status['progress'] = 100
            add_log(f'Current system scan complete: {valid_events} events extracted from {total_events} raw logs', 'success')

        except FileNotFoundError:
            status['state'] = 'error'
            status['message'] = 'journalctl command not found on this system.'
            add_log('Error: journalctl not found. This feature requires a Linux system with systemd.', 'error')
        except Exception as e:
            status['state'] = 'error'
            status['message'] = f'Log scan failed: {str(e)}'
            add_log(f'Log scan failed: {str(e)}', 'error')
        finally:
            status['last_update'] = datetime.now().isoformat()

    # Run in background thread
    thread = threading.Thread(target=scan_task, daemon=True)
    thread.start()

    return jsonify({'message': 'System scan started'})

@app.route('/api/analyze', methods=['POST'])
def analyze_threats():
    """Run threat detection analysis"""
    global status

    data = request.get_json()
    threshold = data.get('threshold', 40)

    def analyze_task():
        global status
        try:
            status['state'] = 'processing'
            status['message'] = 'Building baselines...'
            status['progress'] = 30
            add_log(f'Starting analysis (threshold: {threshold})', 'info')

            alerts = run_detection(risk_threshold=threshold)

            status['progress'] = 100
            status['state'] = 'idle'
            status['message'] = f'Analysis complete - {len(alerts)} threats found'
            add_log(f'Analysis complete: {len(alerts)} threats detected', 'success')

            # Store results for retrieval
            with open('latest_analysis.json', 'w') as f:
                json.dump(alerts, f, indent=2)

        except Exception as e:
            status['state'] = 'error'
            status['message'] = f'Analysis failed: {str(e)}'
            add_log(f'Analysis failed: {str(e)}', 'error')
        finally:
            status['last_update'] = datetime.now().isoformat()

    thread = threading.Thread(target=analyze_task, daemon=True)
    thread.start()

    return jsonify({'message': 'Analysis started'})

@app.route('/api/results', methods=['GET'])
def get_results():
    """Get latest analysis results"""
    try:
        if os.path.exists('latest_analysis.json'):
            with open('latest_analysis.json', 'r') as f:
                alerts = json.load(f)

            # Calculate statistics
            if alerts:
                risk_breakdown = {
                    'CRITICAL': len([a for a in alerts if a['risk_level'] == 'CRITICAL']),
                    'HIGH': len([a for a in alerts if a['risk_level'] == 'HIGH']),
                    'MEDIUM': len([a for a in alerts if a['risk_level'] == 'MEDIUM']),
                    'LOW': len([a for a in alerts if a['risk_level'] == 'LOW'])
                }
            else:
                risk_breakdown = {'CRITICAL': 0, 'HIGH': 0, 'MEDIUM': 0, 'LOW': 0}

            return jsonify({
                'alerts': alerts,
                'total': len(alerts),
                'risk_breakdown': risk_breakdown,
                'timestamp': datetime.now().isoformat()
            })
        else:
            return jsonify({
                'alerts': [],
                'total': 0,
                'risk_breakdown': {'CRITICAL': 0, 'HIGH': 0, 'MEDIUM': 0, 'LOW': 0},
                'timestamp': datetime.now().isoformat()
            })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/events', methods=['GET'])
def get_events():
    """Get paginated event list with search/filter"""
    try:
        page = request.args.get('page', 1, type=int)
        per_page = request.args.get('per_page', 50, type=int)
        search = request.args.get('search', '').strip()
        event_type = request.args.get('event_type', '').strip()
        success_filter = request.args.get('success', '').strip()

        conn = sqlite3.connect(DB_NAME)
        conn.row_factory = sqlite3.Row

        # Build query with filters
        where_clauses: list[str] = []
        params: list[object] = []

        if search:
            where_clauses.append("(username LIKE ? OR source_ip LIKE ?)")
            params.extend([f'%{search}%', f'%{search}%'])
        if event_type:
            where_clauses.append("event_type = ?")
            params.append(event_type)
        if success_filter != '':
            where_clauses.append("success = ?")
            params.append(str(success_filter))

        where_sql = " WHERE " + " AND ".join(where_clauses) if where_clauses else ""

        # Get total count
        count_query = f"SELECT COUNT(*) FROM events{where_sql}"
        total = conn.execute(count_query, params).fetchone()[0]

        # Get paginated results
        offset = (page - 1) * per_page
        data_query = f"SELECT * FROM events{where_sql} ORDER BY id DESC LIMIT ? OFFSET ?"
        rows = conn.execute(data_query, params + [per_page, offset]).fetchall()

        # Get distinct event types for filter dropdown
        event_types = [r[0] for r in conn.execute("SELECT DISTINCT event_type FROM events ORDER BY event_type").fetchall()]

        conn.close()

        events = [dict(row) for row in rows]

        return jsonify({
            'events': events,
            'total': total,
            'page': page,
            'per_page': per_page,
            'total_pages': max(1, (total + per_page - 1) // per_page),
            'event_types': event_types
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/user/<username>', methods=['GET'])
def get_user(username):
    """Get detailed information about a specific user"""
    try:
        details = get_user_details(username)
        if details is None:
            return jsonify({'error': 'User not found'}), 404

        # Convert any pandas Timestamps to strings
        if 'baseline' in details:
            for key, val in details['baseline'].items():
                if hasattr(val, 'isoformat'):
                    details['baseline'][key] = val.isoformat()  # pyre-ignore
                elif isinstance(val, dict):
                    details['baseline'][key] = {str(k): v for k, v in val.items()}

        if 'recent_events' in details:
            for event in details['recent_events']:
                for key, val in event.items():
                    if hasattr(val, 'isoformat'):
                        event[key] = val.isoformat()  # pyre-ignore

        return jsonify(details)
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/logs', methods=['GET'])
def get_logs():
    """Get activity log entries"""
    since = request.args.get('since', 0, type=int)
    if since < len(activity_log):
        return jsonify({
            'logs': list(activity_log)[since:],  # pyre-ignore
            'total': len(activity_log)
        })
    return jsonify({'logs': [], 'total': len(activity_log)})

@app.route('/api/clear', methods=['POST'])
def clear_database():
    """Clear all database data"""
    global status
    try:
        clear_db()

        # Remove cached results
        if os.path.exists('latest_analysis.json'):
            os.remove('latest_analysis.json')

        status['message'] = 'Database cleared'
        status['state'] = 'idle'
        status['last_update'] = datetime.now().isoformat()
        add_log('Database cleared', 'warning')

        return jsonify({'message': 'Database cleared successfully'})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/export', methods=['POST'])
def export_data():
    """Export analysis results in JSON or CSV"""
    try:
        data = request.get_json()
        format_type = data.get('format', 'json')

        if not os.path.exists('latest_analysis.json'):
            return jsonify({'error': 'No analysis results available'}), 404

        with open('latest_analysis.json', 'r') as f:
            alerts = json.load(f)

        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')

        if format_type == 'csv':
            output = io.StringIO()
            if alerts:
                fieldnames = ['user', 'risk', 'risk_level', 'total_events',
                              'failed_events', 'unique_ips', 'first_seen', 'last_seen']
                writer = csv.DictWriter(output, fieldnames=fieldnames)
                writer.writeheader()
                for alert in alerts:
                    row: dict[str, object] = {str(k): v for k, v in alert.items() if k in fieldnames}
                    writer.writerow(row)

            add_log(f'Exported CSV report ({len(alerts)} alerts)', 'success')
            return Response(
                output.getvalue(),
                mimetype='text/csv',
                headers={'Content-Disposition': f'attachment; filename=threat_report_{timestamp}.csv'}
            )
        else:
            filename = f'threat_report_{timestamp}.json'
            with open(filename, 'w') as f:
                json.dump(alerts, f, indent=2)
            add_log(f'Exported JSON report ({len(alerts)} alerts)', 'success')
            return send_file(filename, as_attachment=True)

    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    print("=" * 60)
    print("🔒 IdentityAI Web Interface")
    print("=" * 60)
    print("🌐 Server starting at: http://localhost:5000")
    print("💡 Open this URL in your browser")
    print("🎨 Features: Dark/Light themes, Real-time updates, Animations")
    print("=" * 60)
    app.run(debug=True, host='0.0.0.0', port=5000, threaded=True)