sameersyed/Defence_FrameWork

1import sqlite32from werkzeug.security import generate_password_hash, check_password_hash3from datetime import datetime4 5class Database:6    def __init__(self, db_name='security_app.db'):7        self.db_name = db_name8        self.init_db()9    10    def get_connection(self):11        return sqlite3.connect(self.db_name)12    13    def init_db(self):14        conn = self.get_connection()15        cursor = conn.cursor()16        17        # Users table18        cursor.execute('''19            CREATE TABLE IF NOT EXISTS users (20                id INTEGER PRIMARY KEY AUTOINCREMENT,21                username TEXT UNIQUE NOT NULL,22                email TEXT UNIQUE NOT NULL,23                password_hash TEXT NOT NULL,24                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP25            )26        ''')27        28        # Chat history table29        cursor.execute('''30            CREATE TABLE IF NOT EXISTS chat_history (31                id INTEGER PRIMARY KEY AUTOINCREMENT,32                user_id INTEGER,33                message TEXT,34                response TEXT,35                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,36                FOREIGN KEY (user_id) REFERENCES users (id)37            )38        ''')39        40        # Threat history table41        cursor.execute('''42            CREATE TABLE IF NOT EXISTS threat_history (43                id INTEGER PRIMARY KEY AUTOINCREMENT,44                user_id INTEGER NOT NULL,45                threat_type TEXT NOT NULL,46                risk_level TEXT NOT NULL,47                model_type TEXT,48                confidence INTEGER,49                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,50                FOREIGN KEY (user_id) REFERENCES users(id)51            )52        ''')53        54        conn.commit()55        conn.close()56    57    def create_user(self, username, email, password):58        conn = self.get_connection()59        cursor = conn.cursor()60        61        try:62            password_hash = generate_password_hash(password)63            cursor.execute(64                'INSERT INTO users (username, email, password_hash) VALUES (?, ?, ?)',65                (username, email, password_hash)66            )67            conn.commit()68            return True69        except sqlite3.IntegrityError:70            return False71        finally:72            conn.close()73    74    def verify_user(self, username, password):75        conn = self.get_connection()76        cursor = conn.cursor()77        78        cursor.execute(79            'SELECT id, password_hash FROM users WHERE username = ? OR email = ?',80            (username, username)81        )82        user = cursor.fetchone()83        conn.close()84        85        if user and check_password_hash(user[1], password):86            return user[0]  # Return user ID87        return None88    89    def get_user_by_id(self, user_id):90        conn = self.get_connection()91        cursor = conn.cursor()92        93        cursor.execute('SELECT username, email FROM users WHERE id = ?', (user_id,))94        user = cursor.fetchone()95        conn.close()96        97        return user98    99    def save_chat(self, user_id, message, response):100        conn = self.get_connection()101        cursor = conn.cursor()102        103        cursor.execute(104            'INSERT INTO chat_history (user_id, message, response) VALUES (?, ?, ?)',105            (user_id, message, response)106        )107        conn.commit()108        conn.close()109    110    def save_threat(self, user_id, threat_type, risk_level, model_type='general', confidence=75):111        conn = self.get_connection()112        cursor = conn.cursor()113        114        cursor.execute(115            'INSERT INTO threat_history (user_id, threat_type, risk_level, model_type, confidence) VALUES (?, ?, ?, ?, ?)',116            (user_id, threat_type, risk_level, model_type, confidence)117        )118        conn.commit()119        conn.close()120    121    def get_threat_history(self, user_id, limit=10):122        conn = self.get_connection()123        cursor = conn.cursor()124        125        cursor.execute(126            'SELECT threat_type, risk_level, model_type, confidence, timestamp FROM threat_history WHERE user_id = ? ORDER BY timestamp DESC LIMIT ?',127            (user_id, limit)128        )129        threats = cursor.fetchall()130        conn.close()131        132        return threats133