import os import re import time import logging import threading from flask import Flask, render_template, request, jsonify, redirect, url_for from flask_sqlalchemy import SQLAlchemy from werkzeug.utils import secure_filename from sqlalchemy import create_engine, text, inspect from huggingface_hub import HfApi, hf_hub_download # Set up logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) app = Flask(__name__) # --- App Settings --- UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'uploads') app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER app.config['MAX_CONTENT_LENGTH'] = 32 * 1024 * 1024 # 32MB upload limit os.makedirs(UPLOAD_FOLDER, exist_ok=True) # --- Hugging Face Dataset Persistence Config --- HF_TOKEN = os.environ.get("HF_TOKEN") DATASET_REPO_ID = os.environ.get("DATASET_REPO_ID") if DATASET_REPO_ID: DATASET_REPO_ID = DATASET_REPO_ID.replace("https://huggingface.co/datasets/", "") DATASET_REPO_ID = DATASET_REPO_ID.replace("https://huggingface.co/", "") DATASET_REPO_ID = DATASET_REPO_ID.strip("/") DB_FILENAME = "app_records.db" DB_FILE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), DB_FILENAME) # 1. Download database from Hugging Face Dataset on startup if HF_TOKEN and DATASET_REPO_ID: logger.info("Attempting to download database from Hugging Face Dataset repo...") try: hf_hub_download( repo_id=DATASET_REPO_ID, filename=DB_FILENAME, repo_type="dataset", token=HF_TOKEN, local_dir=os.path.dirname(os.path.abspath(__file__)) ) logger.info("Database downloaded successfully.") except Exception as e: logger.warning(f"Could not download database (might be first run or empty dataset): {e}") else: logger.warning("HF_TOKEN or DATASET_REPO_ID missing. Starting with a fresh database locally.") # 2. Configure Flask-SQLAlchemy to use the SQLite file app.config['SQLALCHEMY_DATABASE_URI'] = f"sqlite:///{DB_FILE_PATH}" app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False db = SQLAlchemy(app) # --- Database Schema Models --- class CaseRecord(db.Model): __tablename__ = 'case_records' id = db.Column(db.Integer, primary_key=True) case_name = db.Column(db.String(100), nullable=False) description = db.Column(db.Text, nullable=True) created_at = db.Column(db.Float, default=time.time) # Cascade delete files when a case is deleted files = db.relationship('EvidenceFile', backref='case', cascade='all, delete-orphan', lazy=True) def __init__(self, case_name: str, description: str = None): self.case_name = case_name self.description = description self.created_at = time.time() class EvidenceFile(db.Model): __tablename__ = 'evidence_files' id = db.Column(db.Integer, primary_key=True) case_id = db.Column(db.Integer, db.ForeignKey('case_records.id', ondelete='CASCADE'), nullable=True) filename = db.Column(db.String(255), nullable=False) size_bytes = db.Column(db.Integer, nullable=False) created_at = db.Column(db.Float, default=time.time) def __init__(self, filename: str, size_bytes: int, case_id: int = None): self.filename = filename self.size_bytes = size_bytes self.case_id = case_id self.created_at = time.time() # 3. Synchronize database back to Hugging Face Dataset in the background def sync_db_to_hf(): if HF_TOKEN and DATASET_REPO_ID: try: logger.info("Syncing database to Hugging Face Dataset...") api = HfApi() api.upload_file( path_or_fileobj=DB_FILE_PATH, path_in_repo=DB_FILENAME, repo_id=DATASET_REPO_ID, repo_type="dataset", token=HF_TOKEN ) logger.info("Database synced to Hugging Face successfully.") except Exception as e: logger.error(f"Error syncing database to Hugging Face Dataset: {e}") # Initialize database schema tables inside application context try: with app.app_context(): inspector = inspect(db.engine) if not inspector.has_table("case_records") or not inspector.has_table("evidence_files"): db.create_all() logger.info("Database schemas initialized successfully.") else: logger.info("Database schemas already exist. Skipping creation.") except Exception as e: logger.warning(f"Database table verification/creation skipped: {e}") # --- Routes --- @app.route('/') def home(): try: cases = CaseRecord.query.order_by(CaseRecord.created_at.desc()).all() files = EvidenceFile.query.order_by(EvidenceFile.created_at.desc()).all() db_status = "HF Dataset Connected" if HF_TOKEN and DATASET_REPO_ID else "Local SQLite Mode" except Exception as e: logger.error(f"Error reading from database: {e}") cases, files = [], [] db_status = "Disconnected (Offline Mode)" return render_template('index.html', cases=cases, files=files, db_status=db_status) @app.route('/submit', methods=['POST']) def submit_form(): case_name = request.form.get('case_name', '').strip() description = request.form.get('description', '').strip() # Validations if not case_name or len(case_name) < 3 or len(case_name) > 100: return jsonify({"success": False, "error": "Case name must be between 3 and 100 characters."}), 400 try: new_record = CaseRecord(case_name=case_name, description=description) db.session.add(new_record) db.session.commit() # Trigger background sync to Hugging Face Dataset threading.Thread(target=sync_db_to_hf, daemon=True).start() return jsonify({"success": True, "message": "Case successfully logged!"}) except Exception as e: logger.error(f"Failed to submit database entry: {e}") db.session.rollback() return jsonify({"success": False, "error": "Database error occurred."}), 500 @app.route('/delete-case/', methods=['POST']) def delete_case(case_id): try: case = CaseRecord.query.get(case_id) if not case: return jsonify({"success": False, "error": "Case not found."}), 404 # Delete files associated from disk for file_record in case.files: file_path = os.path.join(app.config['UPLOAD_FOLDER'], file_record.filename) if os.path.exists(file_path): os.remove(file_path) db.session.delete(case) db.session.commit() # Trigger background sync threading.Thread(target=sync_db_to_hf, daemon=True).start() return jsonify({"success": True, "message": "Case and associated files deleted."}) except Exception as e: logger.error(f"Error deleting case: {e}") db.session.rollback() return jsonify({"success": False, "error": "Failed to delete case."}), 500 @app.route('/upload', methods=['POST']) def upload_file(): if 'file' not in request.files: return jsonify({"success": False, "error": "No file part in the request."}), 400 file = request.files['file'] case_id_val = request.form.get('case_id', '').strip() if file.filename == '': return jsonify({"success": False, "error": "No file selected."}), 400 case_id = int(case_id_val) if case_id_val and case_id_val.isdigit() else None if file: filename = secure_filename(file.filename) save_path = os.path.join(app.config['UPLOAD_FOLDER'], filename) try: file.save(save_path) size_bytes = os.path.getsize(save_path) # Create Database Record new_file = EvidenceFile(filename=filename, size_bytes=size_bytes, case_id=case_id) db.session.add(new_file) db.session.commit() # Trigger background sync threading.Thread(target=sync_db_to_hf, daemon=True).start() return jsonify({ "success": True, "filename": filename, "size_bytes": size_bytes, "message": "File uploaded successfully." }) except Exception as e: logger.error(f"File upload error: {e}") db.session.rollback() return jsonify({"success": False, "error": "Failed to save file on server."}), 500 @app.route('/delete-file/', methods=['POST']) def delete_file(file_id): try: file_record = EvidenceFile.query.get(file_id) if not file_record: return jsonify({"success": False, "error": "File record not found."}), 404 file_path = os.path.join(app.config['UPLOAD_FOLDER'], file_record.filename) if os.path.exists(file_path): os.remove(file_path) db.session.delete(file_record) db.session.commit() # Trigger background sync threading.Thread(target=sync_db_to_hf, daemon=True).start() return jsonify({"success": True, "message": "Evidence file deleted."}) except Exception as e: logger.error(f"Error deleting file: {e}") db.session.rollback() return jsonify({"success": False, "error": "Failed to delete file."}), 500 @app.route('/health') def health(): try: db.session.execute(text('SELECT 1')) db_healthy = True except Exception as e: logger.error(f"Health check failed for database: {e}") db_healthy = False return jsonify({ "status": "healthy", "database": "connected" if db_healthy else "disconnected", "storage": "writable" if os.access(UPLOAD_FOLDER, os.W_OK) else "readonly", "sync_enabled": bool(HF_TOKEN and DATASET_REPO_ID) }) if __name__ == '__main__': # Run local server app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 7860)), debug=True)