Update app.py

app.py

@@ -1,13 +1,12 @@
 import os
 import json
-import time
 import hashlib
 from datetime import datetime
 import gradio as gr
 from huggingface_hub import InferenceClient
 
 # ============================================================
-# CORE CONFIGURATION
+# CONFIG
 # ============================================================
 
 OMNISCIENT_ROOT = "/opt/omniscient"
@@ -21,14 +20,10 @@ MODEL_OPTIONS = {
 
 FALLBACK_MODEL = "HuggingFaceH4/zephyr-7b-beta"
 
-MAX_HISTORY_PAIRS = 4
-MAX_INPUT_CHARS = 4000
-
-# Ensure directories exist
 os.makedirs(EVIDENCE_ROOT, exist_ok=True)
 
 # ============================================================
-# AUDIT LOGGING (APPEND ONLY)
+# AUDIT LOGGING
 # ============================================================
 
 def append_audit_log(action, details):
@@ -50,7 +45,7 @@ def append_audit_log(action, details):
         json.dump(data, f, indent=2)
 
 # ============================================================
-# HASHING + MANIFEST
+# HASHING
 # ============================================================
 
 def sha256_file(path):
@@ -60,10 +55,13 @@ def sha256_file(path):
             h.update(chunk)
     return h.hexdigest()
 
+# ============================================================
+# EVIDENCE
+# ============================================================
 
 def register_evidence(file):
     if not file:
-        return "No file provided."
+        return "⚠️ No file uploaded."
 
     dest_path = os.path.join(EVIDENCE_ROOT, os.path.basename(file.name))
 
@@ -72,35 +70,33 @@ def register_evidence(file):
 
     file_hash = sha256_file(dest_path)
 
-    manifest_entry = {
-        "file": dest_path,
-        "sha256": file_hash,
-        "registered_utc": datetime.utcnow().isoformat()
-    }
-
     manifest_path = os.path.join(EVIDENCE_ROOT, "manifest.json")
+    manifest = []
 
     if os.path.exists(manifest_path):
         with open(manifest_path, "r") as f:
             manifest = json.load(f)
-    else:
-        manifest = []
 
-    manifest.append(manifest_entry)
+    entry = {
+        "file": dest_path,
+        "sha256": file_hash,
+        "registered_utc": datetime.utcnow().isoformat()
+    }
+
+    manifest.append(entry)
 
     with open(manifest_path, "w") as f:
         json.dump(manifest, f, indent=2)
 
-    append_audit_log("EVIDENCE_REGISTERED", manifest_entry)
-
-    return f"Evidence Registered.\nSHA256: {file_hash}"
+    append_audit_log("EVIDENCE_REGISTERED", entry)
 
+    return f"✅ Evidence Registered\n\nFile: {dest_path}\nSHA256: {file_hash}"
 
 def verify_manifest():
     manifest_path = os.path.join(EVIDENCE_ROOT, "manifest.json")
 
     if not os.path.exists(manifest_path):
-        return "No manifest found."
+        return "⚠️ No manifest found."
 
     with open(manifest_path, "r") as f:
         manifest = json.load(f)
@@ -109,25 +105,25 @@ def verify_manifest():
 
     for entry in manifest:
         path = entry["file"]
-        expected_hash = entry["sha256"]
+        expected = entry["sha256"]
 
         if not os.path.exists(path):
-            results.append(f"MISSING: {path}")
+            results.append(f"❌ MISSING: {path}")
             continue
 
-        actual_hash = sha256_file(path)
+        actual = sha256_file(path)
 
-        if actual_hash == expected_hash:
-            results.append(f"VALID: {path}")
+        if actual == expected:
+            results.append(f"✅ VALID: {path}")
         else:
-            results.append(f"HASH MISMATCH: {path}")
+            results.append(f"⚠️ HASH MISMATCH: {path}")
 
-    append_audit_log("MANIFEST_VERIFICATION", {"result_count": len(results)})
+    append_audit_log("MANIFEST_VERIFIED", {"entries": len(results)})
 
     return "\n".join(results)
 
 # ============================================================
-# AI AUGMENTATION LAYER (READ-ONLY)
+# AI LAYER
 # ============================================================
 
 def get_client(model_id):
@@ -136,7 +132,6 @@ def get_client(model_id):
         raise RuntimeError("HF_TOKEN not set.")
     return InferenceClient(model=model_id, token=token)
 
-
 def run_ai(model_id, messages, temperature=0.4):
     try:
         client = get_client(model_id)
@@ -148,8 +143,8 @@ def run_ai(model_id, messages, temperature=0.4):
         )
         return response.choices[0].message.content
     except Exception:
-        fallback_client = get_client(FALLBACK_MODEL)
-        response = fallback_client.chat_completion(
+        fallback = get_client(FALLBACK_MODEL)
+        response = fallback.chat_completion(
             messages=messages,
             max_tokens=900,
             temperature=temperature,
@@ -157,120 +152,63 @@ def run_ai(model_id, messages, temperature=0.4):
         )
         return "⚠️ Fallback model used.\n\n" + response.choices[0].message.content
 
-
-def ai_augment_analysis(message, model_id):
+def ai_analysis(input_text, model_label):
     system_prompt = """
-You are an AI analytical augmentation module for digital forensics.
-
-You may:
-- Summarize evidence
-- Identify technical patterns
-- Suggest hypotheses
-- Outline investigative steps
-
-You may NOT:
-- Modify evidence
-- Make legal determinations
-- Assert conclusions beyond observed data
-
-Structure output as:
+You are an AI augmentation module for digital forensics.
 
+Structure:
 1. Observed Technical Artifacts
 2. Analytical Interpretation
 3. Competing Hypotheses
 4. Risk-Relevant Observations
 5. Limitations
 6. Recommended Human Review Steps
+
+AI is advisory only.
 """
 
     messages = [
         {"role": "system", "content": system_prompt},
-        {"role": "user", "content": message}
+        {"role": "user", "content": input_text}
     ]
 
-    result = run_ai(model_id, messages)
-
-    append_audit_log("AI_ANALYSIS", {"input_length": len(message)})
+    result = run_ai(MODEL_OPTIONS[model_label], messages)
+    append_audit_log("AI_ANALYSIS", {"length": len(input_text)})
 
     return result
 
 # ============================================================
-# DOCUMENTATION ENGINE
+# DOCUMENTATION
 # ============================================================
 
-def generate_system_documentation(model_id):
-    prompt = """
-Generate formal system documentation for Omniscient Investigative Infrastructure.
-
-Include:
-- Executive Overview
-- Architecture
-- Evidence Handling Model
-- Hashing & Integrity Controls
-- AI Governance Model
-- Audit Logging Controls
-- Court Defensibility Posture
-- Operational Boundaries
-
-Write in professional enterprise documentation format.
-"""
+def generate_documentation(doc_type, model_label):
 
-    messages = [
-        {"role": "system", "content": "You are a senior DFIR infrastructure documentation specialist."},
-        {"role": "user", "content": prompt}
-    ]
-
-    return run_ai(model_id, messages)
-
-
-def generate_dfir_practice_documentation(model_id):
-    prompt = """
-Generate professional digital forensics best practice documentation.
+    if doc_type == "System Architecture":
+        prompt = "Generate formal Omniscient system architecture documentation."
 
-Include:
-- Evidence Preservation
-- Imaging Standards
-- Hash Verification
-- Logging & Audit Trails
-- AI Augmentation Boundaries
-- Legal Safeguards
-- Expert Witness Preparation
+    elif doc_type == "DFIR Best Practices":
+        prompt = "Generate professional digital forensics best practice documentation."
 
-Write formally.
-"""
+    else:
+        return "Invalid documentation type."
 
     messages = [
-        {"role": "system", "content": "You are a senior DFIR practitioner."},
+        {"role": "system", "content": "You are a professional DFIR documentation specialist."},
         {"role": "user", "content": prompt}
     ]
 
-    return run_ai(model_id, messages)
+    return run_ai(MODEL_OPTIONS[model_label], messages)
 
 # ============================================================
-# ROUTER
+# SYSTEM STATUS
 # ============================================================
 
-def manager_router(message, model_label, mode, file_upload):
-
-    model_id = MODEL_OPTIONS[model_label]
-    message = (message or "")[:MAX_INPUT_CHARS]
-
-    if mode == "Register Evidence":
-        return register_evidence(file_upload)
-
-    if mode == "Verify Manifest Integrity":
-        return verify_manifest()
-
-    if mode == "AI Augmentation Analysis":
-        return ai_augment_analysis(message, model_id)
-
-    if mode == "Generate System Documentation":
-        return generate_system_documentation(model_id)
-
-    if mode == "Generate DFIR Practice Documentation":
-        return generate_dfir_practice_documentation(model_id)
-
-    return "Invalid mode selected."
+def system_status():
+    return f"""
+Omniscient Root: {OMNISCIENT_ROOT}
+Evidence Directory Exists: {os.path.exists(EVIDENCE_ROOT)}
+Audit Log Exists: {os.path.exists(LOG_FILE)}
+"""
 
 # ============================================================
 # UI
@@ -278,54 +216,52 @@ def manager_router(message, model_label, mode, file_upload):
 
 with gr.Blocks(title="Omniscient Investigative Infrastructure") as demo:
 
-    gr.Markdown("## Omniscient — Investigative Infrastructure & DFIR AI Augmentation")
-
-    with gr.Row():
-
-        with gr.Column(scale=3):
+    gr.Markdown("## Omniscient — Investigative Infrastructure Control Panel")
 
-            output_box = gr.Textbox(lines=30, label="Output")
-
-            msg = gr.Textbox(
-                placeholder="Enter analysis request or documentation command...",
-                label="Input"
-            )
+    with gr.Tabs():
 
+        # Evidence Tab
+        with gr.Tab("Evidence Management"):
             file_upload = gr.File(label="Upload Evidence File")
+            register_btn = gr.Button("Register Evidence")
+            evidence_output = gr.Textbox(lines=10)
 
-            execute_btn = gr.Button("Execute")
-
-        with gr.Column(scale=1):
+        # Integrity Tab
+        with gr.Tab("Integrity Verification"):
+            verify_btn = gr.Button("Verify Manifest Integrity")
+            integrity_output = gr.Textbox(lines=15)
 
+        # AI Tab
+        with gr.Tab("AI Augmentation"):
+            ai_input = gr.Textbox(lines=8, label="Analysis Input")
             model_selector = gr.Dropdown(
                 choices=list(MODEL_OPTIONS.keys()),
                 value=list(MODEL_OPTIONS.keys())[0],
                 label="AI Model"
             )
-
-            mode_selector = gr.Dropdown(
-                choices=[
-                    "Register Evidence",
-                    "Verify Manifest Integrity",
-                    "AI Augmentation Analysis",
-                    "Generate System Documentation",
-                    "Generate DFIR Practice Documentation",
-                ],
-                value="AI Augmentation Analysis",
-                label="Operation Mode"
+            ai_btn = gr.Button("Run AI Analysis")
+            ai_output = gr.Textbox(lines=20)
+
+        # Documentation Tab
+        with gr.Tab("Documentation"):
+            doc_type = gr.Dropdown(
+                choices=["System Architecture", "DFIR Best Practices"],
+                value="System Architecture",
+                label="Documentation Type"
             )
-
-    execute_btn.click(
-        manager_router,
-        inputs=[msg, model_selector, mode_selector, file_upload],
-        outputs=output_box,
-    )
-
-    msg.submit(
-        manager_router,
-        inputs=[msg, model_selector, mode_selector, file_upload],
-        outputs=output_box,
-    )
+            doc_btn = gr.Button("Generate Documentation")
+            doc_output = gr.Textbox(lines=20)
+
+        # Status Tab
+        with gr.Tab("System Status"):
+            status_btn = gr.Button("Check System Status")
+            status_output = gr.Textbox(lines=8)
+
+    register_btn.click(register_evidence, inputs=file_upload, outputs=evidence_output)
+    verify_btn.click(verify_manifest, outputs=integrity_output)
+    ai_btn.click(ai_analysis, inputs=[ai_input, model_selector], outputs=ai_output)
+    doc_btn.click(generate_documentation, inputs=[doc_type, model_selector], outputs=doc_output)
+    status_btn.click(system_status, outputs=status_output)
 
 if __name__ == "__main__":
     demo.queue()