Upload 7 files

DockerFile

@@ -0,0 +1,17 @@
+FROM python:3.9
+
+WORKDIR /app
+
+COPY requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY . .
+
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+EXPOSE 7860
+
+CMD ["gunicorn", "-b", "0.0.0.0:7860", "app:app", "--timeout", "120"]

app.py

@@ -9,13 +9,12 @@ import os
 import subprocess
 import json
 import time
+import sys
 from pymongo import MongoClient
 from routes.reviews import reviews_bp
-from flask_limiter import Limiter
-from flask_limiter.util import get_remote_address
+from extensions import limiter
 from datetime import datetime
 from datetime import timedelta
-from bson import ObjectId
 from models.reviews import reviews_collection
 from schemas import ScanRequest
 from pydantic import ValidationError
@@ -26,9 +25,13 @@ load_dotenv()
 
 app = Flask(__name__)
 Compress(app)
-CORS(app)
 
-# Load environment variables
+FRONTEND_URL = os.getenv("FRONTEND_URL", "http://localhost:5173")
+allowed_origins = list({FRONTEND_URL, "http://localhost:5173", "http://localhost:3000", "http://localhost:8080"})
+CORS(app, origins=allowed_origins)
+
+app.config['MAX_CONTENT_LENGTH'] = 1 * 1024 * 1024
+
 env = os.getenv('FLASK_ENV', 'development')
 if env == 'development':
     load_dotenv('.env.development')
@@ -47,9 +50,7 @@ else:
         "CACHE_DEFAULT_TIMEOUT": 3600
     })
 
-cache.init_app(app)
-
-cache.init_app(app)
+cache.init_app(app)  
 
 def files_hash(files):
     h = hashlib.sha256()
@@ -75,10 +76,6 @@ scan_history = db["scan_history"]
 
 app.register_blueprint(auth_bp, url_prefix="/api")
 
-limiter = Limiter(
-    key_func=get_remote_address,
-    default_limits=["20 per minute"]
-)
 limiter.init_app(app)
 
 @app.route('/api/scan', methods=['POST'])
@@ -87,7 +84,6 @@ limiter.init_app(app)
 def scan_code():
     try:
         data = request.get_json()
-        app.logger.info(f"Incoming request: {data}")   
 
         try:
             req = ScanRequest(**data)
@@ -95,11 +91,13 @@ def scan_code():
             app.logger.error(f"Validation error: {e.errors()}")
             return jsonify({"error": e.errors()}), 400
 
-        files = [f.dict() for f in req.files]
+        files = [f.model_dump() for f in req.files]
         language = req.language.lower()
+        app.logger.info(f"Scan request: {len(files)} files, language={language}")
+
         username = get_jwt_identity()
 
-        key = f"scan:{username}:{files_hash(files)}"
+        key = f"scan:{username}:{language}:{files_hash(files)}"
         cached = cache.get(key)
         if cached:
             return jsonify({"result": cached, "cached": True})
@@ -111,22 +109,28 @@ def scan_code():
                     code_file.write(f["content"])
 
             if language == "python":
-                scan_command = ["python", "-m", "bandit", "-r", temp_dir, "-f", "json"]
+                scan_command = [sys.executable, "-m", "bandit", "-r", temp_dir, "-f", "json"]
             elif language == "javascript":
-                scan_command = ["python", "-m", "semgrep", "--config=p/javascript", "--json", temp_dir]
+                scan_command = [sys.executable, "-m", "semgrep", "--config=p/javascript", "--json", temp_dir]
             else:
                 return jsonify({"error": "Unsupported language"}), 400
 
             app.logger.info(f"Running: {' '.join(scan_command)}")
             result = subprocess.run(scan_command, capture_output=True, text=True)
 
-            app.logger.info(f"stdout: {result.stdout[:500]}")
             app.logger.info(f"stderr: {result.stderr}")
 
             if result.returncode not in (0, 1, 2):
                 return jsonify({"error": result.stderr}), 500
+            
+            app.logger.info("Scan completed successfully")
 
             try:
+                if not result.stdout.strip():
+                    return jsonify({
+                        "error": "Scanner produced no output",
+                        "stderr": result.stderr
+                    }), 500
                 output_json = json.loads(result.stdout)
             except Exception as e:
                 return jsonify({"error": f"JSON parse failed: {str(e)}", "raw": result.stdout}), 500
@@ -147,7 +151,6 @@ def scan_code():
         return jsonify({"error": str(e)}), 500
 
 
-
 @app.route("/api/health")
 def health():
     return jsonify({"status": "ok"})
@@ -168,10 +171,8 @@ def enhance():
         if not code.strip():
             return jsonify({"error": "No code provided"}), 400
 
-        # 🔹 New format (returns dict)
         result = enhance_code(code, language)
 
-        # Save to history (with candidates + explanations)
         enhance_history.insert_one({
             "username": username,
             "code": code,
@@ -195,11 +196,9 @@ def history():
     try:
         username = get_jwt_identity()
 
-        # Fetch both histories
         enhance_records = list(enhance_history.find({"username": username}).sort("timestamp", -1))
         scan_records = list(scan_history.find({"username": username}).sort("timestamp", -1))
 
-        # Convert ObjectId to string & return only relevant fields
         def clean(record, record_type):
             return {
                 "id": str(record.get("_id")),
@@ -207,8 +206,8 @@ def history():
                 "code": record.get("code"),
                 "enhanced_code": record.get("enhanced_code"),
                 "diff": record.get("diff"),
-                "candidates": record.get("candidates", []),   # ✅ added
-                "explanations": record.get("explanations", []), # ✅ added
+                "candidates": record.get("candidates", []),
+                "explanations": record.get("explanations", []),
                 "result": record.get("result") if record_type == "scan" else None,
                 "timestamp": record.get("timestamp"),
             }
@@ -257,21 +256,22 @@ def submit_review():
 
     except Exception as e:
         return jsonify({"error": str(e)}), 500
-    
+
 
 @app.route("/api/enhance-stream", methods=["POST"])
+@limiter.limit("5/minute")
 @jwt_required()
 def enhance_stream():
     data = request.get_json()
     code = data.get("code", "")
     language = data.get("language", "python")
+    username = get_jwt_identity()  
 
     if not code.strip():
         return jsonify({"error": "No code"}), 400
 
     def generate():
         try:
-            # 1️⃣ starting
             yield json.dumps({
                 "type": "progress",
                 "progress": 5
@@ -279,16 +279,27 @@ def enhance_stream():
 
             time.sleep(0.5)
 
-            # 2️⃣ preprocessing
             yield json.dumps({
                 "type": "progress",
                 "progress": 20
             }) + "\n"
 
-            # 3️⃣ heavy AI call
             result = enhance_code(code, language)
 
-            # 4️⃣ done
+            try:
+                enhance_history.insert_one({
+                    "username": username,
+                    "code": code,
+                    "language": language,
+                    "enhanced_code": result.get("enhanced_code", ""),
+                    "diff": result.get("diff", []),
+                    "candidates": result.get("candidates", []),
+                    "explanations": result.get("explanations", []),
+                    "timestamp": datetime.utcnow().isoformat()
+                })
+            except Exception:
+                pass  
+
             yield json.dumps({
                 "type": "result",
                 "data": result
@@ -306,6 +317,5 @@ def enhance_stream():
     )
 
 
-
 if __name__ == '__main__':
-    app.run(host="0.0.0.0", port=5000, debug=True)
+    app.run(host="0.0.0.0", port=5000, debug=True)

extensions.py

@@ -0,0 +1,7 @@
+from flask_limiter import Limiter
+from flask_limiter.util import get_remote_address
+
+limiter = Limiter(
+    key_func=get_remote_address,
+    default_limits=["20 per minute"]
+)

model.py

@@ -1,264 +1,337 @@
-import torch
-import difflib
-import re
-from transformers import (
-    AutoTokenizer,
-    AutoModelForSeq2SeqLM,
-    AutoModelForCausalLM
-)
-
-# ----------------------------
-# Performance Settings
-# ----------------------------
-
-torch.set_num_threads(2)
-DEVICE = "cpu"  
-
-# ----------------------------
-# Models and their types
-# ----------------------------
-
-MODEL_CONFIGS = {
-    "Salesforce/codet5-base": "seq2seq",        # CodeT5
-    #"EleutherAI/gpt-neo-1.3B": "causal",        # GPT-Neo (disabled due to free hosting for now; enable on local hosting)
-    "microsoft/CodeGPT-small-py": "causal",     # CodeGPT-small (Python)
-}
-
-# ----------------------------
-# Load tokenizers and models
-# ----------------------------
-
-tokenizers = {}
-models = {}
-
-print("🔹 Loading models...")
-
-for name, mtype in MODEL_CONFIGS.items():
-    print(f"Loading {name} ...")
-
-    tokenizers[name] = AutoTokenizer.from_pretrained(name)
-
-    if mtype == "seq2seq":
-        model = AutoModelForSeq2SeqLM.from_pretrained(name)
-    else:
-        model = AutoModelForCausalLM.from_pretrained(name)
-
-    model.to(DEVICE)
-    model.eval()          
-    models[name] = model
-
-print("✅ All models loaded")
-
-# ----------------------------
-# Rule-based fixes
-# ----------------------------
-
-SECURE_REPLACEMENTS = {
-    "hashlib.md5": ("hashlib.sha256", "MD5 is weak, replaced with SHA-256."),
-    "hashlib.sha1": ("hashlib.sha256", "SHA1 is weak, replaced with SHA-256."),
-    "eval(": ("ast.literal_eval(", "Unsafe eval removed, replaced with safe literal_eval."),
-    "pickle.load(": ("# pickle.load removed", "pickle.load is unsafe, consider json/safe loaders."),
-}
-
-def rule_based_patch(code: str):
-    explanations = []
-    patched = code
-
-    for bad, (good, reason) in SECURE_REPLACEMENTS.items():
-        if bad in patched:
-            patched = patched.replace(bad, good)
-            explanations.append({
-                "change": f"{bad} → {good}",
-                "reason": reason
-            })
-
-    return patched, explanations
-
-# ----------------------------
-# Structure preservation
-# ----------------------------
-
-def preserve_structure(original: str, enhanced: str):
-    final_code = enhanced
-
-    original_imports = [
-        l for l in original.splitlines()
-        if l.strip().startswith(("import ", "from "))
-    ]
-
-    for imp in original_imports:
-        if imp not in final_code:
-            final_code = imp + "\n" + final_code
-
-    original_defs = [
-        l for l in original.splitlines()
-        if l.strip().startswith("def ")
-    ]
-
-    for d in original_defs:
-        if d.split("(")[0] not in final_code:
-            final_code = (
-                d +
-                "\n    # [!] Function body missing, please review\n" +
-                final_code
-            )
-
-    return final_code
-
-# ----------------------------
-# Diff creation
-# ----------------------------
-
-def create_diff(original: str, enhanced: str):
-    diff_lines = difflib.unified_diff(
-        original.splitlines(),
-        enhanced.splitlines(),
-        lineterm=""
-    )
-
-    formatted = []
-
-    for line in diff_lines:
-        if line.startswith("+") and not line.startswith("+++"):
-
-            formatted.append({
-                "type": "add",
-                "content": line[1:]
-            })
-
-        elif line.startswith("-") and not line.startswith("---"):
-
-            formatted.append({
-                "type": "remove",
-                "content": line[1:]
-            })
-
-        elif not line.startswith("@@"):
-
-            formatted.append({
-                "type": "context",
-                "content": line
-            })
-
-    return formatted
-
-# ----------------------------
-# Postprocess output
-# ----------------------------
-
-def postprocess_code(code: str):
-    code = re.sub(r'^"""|"""$', '', code.strip())
-    lines = code.splitlines()
-    return "\n".join(
-        l.replace("\t", "    ").rstrip()
-        for l in lines
-    )
-
-# ----------------------------
-# Run one model
-# ----------------------------
-
-def run_model(model_name, code, language):
-
-    tokenizer = tokenizers[model_name]
-    model = models[model_name]
-    mtype = MODEL_CONFIGS[model_name]
-
-    prompt = f"Fix security issues in this {language} code:\n{code}"
-
-    if mtype == "seq2seq":
-
-        inputs = tokenizer(
-            prompt,
-            return_tensors="pt",
-            truncation=True,
-            max_length=512
-        ).to(DEVICE)
-
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=512,
-            num_beams=4
-        )
-
-    else:
-
-        inputs = tokenizer(
-            prompt,
-            return_tensors="pt",
-            truncation=True,
-            max_length=512
-        ).to(DEVICE)
-
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=256,
-            temperature=0.3,
-            top_p=0.95,
-            do_sample=False
-        )
-
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# ----------------------------
-# Main enhancer
-# ----------------------------
-
-def enhance_code(code: str, language: str):
-
-    with torch.no_grad():
-
-        try:
-            # 1️⃣ Rule-based fixes
-            patched_code, rule_explanations = rule_based_patch(code)
-
-            # 2️⃣ Model ensemble
-            candidates = []
-
-            for m in MODEL_CONFIGS.keys():
-                try:
-                    enhanced = run_model(m, patched_code, language)
-                    enhanced = postprocess_code(enhanced)
-                    enhanced = preserve_structure(code, enhanced)
-
-                    candidates.append({
-                        "model": m,
-                        "code": enhanced
-                    })
-
-                except Exception as e:
-                    candidates.append({
-                        "model": m,
-                        "code": f"# [!] Failed: {str(e)}"
-                    })
-
-            # 3️⃣ Choose longest output as best
-            best = max(candidates, key=lambda c: len(c["code"]))
-
-            diff = create_diff(code, best["code"])
-
-            explanations = rule_explanations + [{
-                "change": "Model ensemble",
-                "reason": "Best candidate selected from multiple models"
-            }]
-
-            return {
-                "enhanced_code": best["code"],
-                "diff": diff,
-                "candidates": candidates[:3],
-                "explanations": explanations
-            }
-
-        except Exception as e:
-
-            fallback = code + f"\n# [!] Enhancer crashed: {str(e)}"
-
-            return {
-                "enhanced_code": fallback,
-                "diff": create_diff(code, fallback),
-                "candidates": [],
-                "explanations": [{
-                    "change": "Error",
-                    "reason": str(e)
-                }]
-            }
+import torch
+import difflib
+import re
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM,
+    AutoModelForCausalLM
+)
+import os
+
+torch.set_num_threads(max(1, os.cpu_count() // 2))
+DEVICE = "cpu"
+
+MAX_CODE_CHARS = 8000
+
+MODEL_CONFIGS = {
+    "Salesforce/codet5-base": "seq2seq",        # CodeT5
+    #"EleutherAI/gpt-neo-1.3B": "causal",        # GPT-Neo (disabled due to free hosting)
+    "microsoft/CodeGPT-small-py": "causal",     # CodeGPT-small (Python)
+}
+
+tokenizers = {}
+models = {}
+
+def load_models():
+    for name, mtype in MODEL_CONFIGS.items():
+        print(f"Loading {name} ...")
+
+        tokenizers[name] = AutoTokenizer.from_pretrained(name, use_fast=False)
+
+        if mtype == "seq2seq":
+            model = AutoModelForSeq2SeqLM.from_pretrained(name)
+        else:
+            model = AutoModelForCausalLM.from_pretrained(name)
+
+        model.to(DEVICE)
+        model.eval()
+        models[name] = model
+
+    print("✅ All models loaded")
+
+print("🔹 Loading models...")
+_werkzeug_parent = (
+    os.environ.get("WERKZEUG_RUN_MAIN") is None
+    and os.environ.get("FLASK_DEBUG", "0") in ("1", "true", "True")
+)
+if not _werkzeug_parent:
+    load_models()
+
+SECURE_REPLACEMENTS = {
+    # Weak hashing
+    "hashlib.md5":          ("hashlib.sha256",        "MD5 is cryptographically broken; replaced with SHA-256."),
+    "hashlib.sha1":         ("hashlib.sha256",        "SHA-1 is deprecated for security use; replaced with SHA-256."),
+    # Dangerous execution
+    "eval(":                ("ast.literal_eval(",     "eval() executes arbitrary code; use ast.literal_eval for safe parsing."),
+    "exec(":                ("# exec() removed —",    "exec() executes arbitrary code strings; remove or sandbox."),
+    # Insecure deserialization
+    "pickle.load(":         ("# pickle.load UNSAFE —","pickle.load deserialises arbitrary objects; use json or safer alternatives."),
+    "pickle.loads(":        ("# pickle.loads UNSAFE —","pickle.loads is an RCE risk; use json.loads instead."),
+    "yaml.load(":           ("yaml.safe_load(",       "yaml.load with arbitrary loader executes code; use yaml.safe_load."),
+    # Command injection
+    "os.system(":           ("subprocess.run([",      "os.system passes args to the shell; use subprocess.run with a list to avoid injection."),
+    "shell=True":           ("shell=False",           "shell=True enables command injection; pass args as a list with shell=False."),
+    # Insecure temp files
+    "tempfile.mktemp(":     ("tempfile.mkstemp(",     "mktemp() has a race condition; use mkstemp() which atomically creates the file."),
+    # Weak randomness
+    "random.random()":      ("secrets.token_bytes(16)","random is not cryptographically secure; use the secrets module for security-sensitive values."),
+    "random.randint(":      ("secrets.randbelow(",    "random.randint is not cryptographically secure; use secrets.randbelow for security tokens."),
+    # Insecure TLS
+    "verify=False":         ("verify=True",           "Disabling TLS verification allows MITM attacks; always verify certificates."),
+    "ssl.CERT_NONE":        ("ssl.CERT_REQUIRED",     "ssl.CERT_NONE disables certificate validation entirely; use ssl.CERT_REQUIRED."),
+    # Debug/info leakage
+    "DEBUG = True":         ("DEBUG = False",         "Debug mode exposes stack traces and internal config; disable in production."),
+    "app.run(debug=True":   ("app.run(debug=False",   "Flask debug=True enables the Werkzeug debugger which allows arbitrary code execution."),
+}
+
+PYTHON_EXTRA_REPLACEMENTS = {
+    "% username":           ("# Use parameterised query","String-formatted SQL allows injection; use parameterised queries with ? placeholders."),
+    "format(username":      ("# Use parameterised query","String-formatted SQL allows injection; use parameterised queries."),
+    "http://":              ("https://",              "Unencrypted HTTP transmits data in plaintext; upgrade to HTTPS."),
+}
+
+JS_SECURE_REPLACEMENTS = {
+    "innerHTML =":              ("textContent =",             "innerHTML enables XSS; use textContent to safely set plain text."),
+    "innerHTML+=":              ("textContent+=",             "innerHTML enables XSS; use textContent instead."),
+    "document.write(":          ("// document.write removed —","document.write allows XSS injection; use DOM APIs instead."),
+    "eval(":                    ("JSON.parse(",               "eval() executes arbitrary JavaScript; use JSON.parse for data or a safe alternative."),
+    "Math.random()":            ("crypto.getRandomValues(",   "Math.random is not cryptographically secure; use crypto.getRandomValues."),
+    "http://":                  ("https://",                  "Unencrypted HTTP in JS code; upgrade to HTTPS."),
+    "dangerouslySetInnerHTML":  ("// dangerouslySetInnerHTML — review needed","dangerouslySetInnerHTML bypasses React's XSS protection; sanitise input with DOMPurify first."),
+    "localStorage.setItem":     ("// Consider sessionStorage —","localStorage persists indefinitely; prefer sessionStorage for sensitive session data."),
+}
+
+SECRET_PATTERNS = [
+    (
+        r'(?i)(?:password|passwd|pwd)\s*=\s*["\'][^"\']{8,}["\']',
+        "Hardcoded password detected — move to an environment variable."
+    ),
+    (
+        r'(?i)(?:api_key|apikey|secret_key|secret|auth_token)\s*=\s*["\'][a-zA-Z0-9+/=_\-]{16,}["\']',
+        "Hardcoded API key or secret detected — move to an environment variable."
+    ),
+    (
+        r'(?:AKIA|ASIA)[A-Z0-9]{16}',
+        "AWS Access Key ID pattern detected — never hardcode AWS credentials."
+    ),
+    (
+        r'(?i)private_key\s*=\s*["\'][^"\']{10,}["\']',
+        "Hardcoded private key detected — load from a secure vault or environment variable."
+    ),
+]
+
+def scan_secrets(code: str) -> list:
+    """Detect hardcoded secrets via regex and return explanation entries."""
+    findings = []
+    for pattern, reason in SECRET_PATTERNS:
+        for match in re.finditer(pattern, code):
+            snippet = match.group()[:60]
+            findings.append({
+                "change": f"Hardcoded secret: {snippet}{'...' if len(match.group()) > 60 else ''}",
+                "reason": reason
+            })
+    return findings
+
+def rule_based_patch(code: str, language: str = "python"):
+    explanations = []
+    patched = code
+
+    for bad, (good, reason) in SECURE_REPLACEMENTS.items():
+        if bad in patched:
+            patched = patched.replace(bad, good)
+            explanations.append({
+                "change": f"{bad} → {good}",
+                "reason": reason
+            })
+
+    lang_rules = JS_SECURE_REPLACEMENTS if language == "javascript" else PYTHON_EXTRA_REPLACEMENTS
+    for bad, (good, reason) in lang_rules.items():
+        if bad in patched:
+            patched = patched.replace(bad, good)
+            explanations.append({
+                "change": f"{bad} → {good}",
+                "reason": reason
+            })
+
+    secret_findings = scan_secrets(code)
+    explanations.extend(secret_findings)
+
+    return patched, explanations
+
+def preserve_structure(original: str, enhanced: str):
+    final_code = enhanced
+
+    original_imports = [
+        l for l in original.splitlines()
+        if l.strip().startswith(("import ", "from "))
+    ]
+
+    for imp in original_imports:
+        if imp not in final_code:
+            final_code = imp + "\n" + final_code
+
+    original_defs = [
+        l for l in original.splitlines()
+        if l.strip().startswith("def ")
+    ]
+
+    for d in original_defs:
+        if d.split("(")[0] not in final_code:
+            final_code = (
+                d +
+                "\n    # [!] Function body missing, please review\n" +
+                final_code
+            )
+
+    return final_code
+
+def create_diff(original: str, enhanced: str):
+    diff_lines = difflib.unified_diff(
+        original.splitlines(),
+        enhanced.splitlines(),
+        lineterm=""
+    )
+
+    formatted = []
+
+    for line in diff_lines:
+        if line.startswith("+") and not line.startswith("+++"):
+            formatted.append({
+                "type": "add",
+                "content": line[1:]
+            })
+        elif line.startswith("-") and not line.startswith("---"):
+            formatted.append({
+                "type": "remove",
+                "content": line[1:]
+            })
+        elif not line.startswith("@@"):
+            formatted.append({
+                "type": "context",
+                "content": line
+            })
+
+    return formatted
+
+def postprocess_code(code: str):
+    code = re.sub(r'^"""|"""$', '', code.strip())
+    lines = code.splitlines()
+    return "\n".join(
+        l.replace("\t", "    ").rstrip()
+        for l in lines
+    )
+
+def score_candidate(candidate_code: str, original_code: str) -> int:
+    """
+    Score a candidate by how many known bad patterns it fixed
+    minus any new bad patterns it introduced.
+    Failed/crashed candidates are heavily penalised.
+    """
+    if "# [!] Failed" in candidate_code[:80]:
+        return -9999
+
+    all_bad = (
+        list(SECURE_REPLACEMENTS.keys()) +
+        list(PYTHON_EXTRA_REPLACEMENTS.keys()) +
+        list(JS_SECURE_REPLACEMENTS.keys())
+    )
+    fixed = sum(1 for p in all_bad if p in original_code and p not in candidate_code)
+    new_issues = sum(1 for p in all_bad if p not in original_code and p in candidate_code)
+    return fixed - new_issues
+
+def run_model(model_name, code, language):
+    tokenizer = tokenizers[model_name]
+    model = models[model_name]
+    mtype = MODEL_CONFIGS[model_name]
+
+    prompt = f"Fix security issues in this {language} code:\n{code}"
+
+    if mtype == "seq2seq":
+        inputs = tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=512
+        ).to(DEVICE)
+
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=512,
+            num_beams=4
+        )
+    else:
+        inputs = tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=512
+        ).to(DEVICE)
+
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            temperature=0.3,
+            top_p=0.95,
+            do_sample=True
+        )
+
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+def enhance_code(code: str, language: str):
+
+    if len(code) > MAX_CODE_CHARS:
+        return {
+            "enhanced_code": code,
+            "diff": [],
+            "candidates": [],
+            "explanations": [{
+                "change": "Input too large",
+                "reason": f"Code exceeds {MAX_CODE_CHARS} character limit. Please split into smaller files."
+            }]
+        }
+
+    with torch.no_grad():
+        try:
+            patched_code, rule_explanations = rule_based_patch(code, language)
+
+            candidates = []
+
+            for m in MODEL_CONFIGS.keys():
+                try:
+                    enhanced = run_model(m, patched_code, language)
+                    enhanced = postprocess_code(enhanced)
+                    enhanced = preserve_structure(code, enhanced)
+
+                    candidates.append({
+                        "model": m,
+                        "code": enhanced
+                    })
+
+                except Exception as e:
+                    candidates.append({
+                        "model": m,
+                        "code": f"# [!] Failed: {str(e)}"
+                    })
+
+            valid_candidates = [c for c in candidates if "# [!] Failed" not in c["code"][:80]]
+            if valid_candidates:
+                best = max(valid_candidates, key=lambda c: score_candidate(c["code"], code))
+            else:
+                best = {"model": "rule-based", "code": patched_code}
+
+            diff = create_diff(code, best["code"])
+
+            explanations = rule_explanations + [{
+                "change": "Model ensemble",
+                "reason": "Best candidate selected from multiple models based on security improvement score"
+            }]
+
+            return {
+                "enhanced_code": best["code"],
+                "diff": diff,
+                "candidates": candidates[:3],
+                "explanations": explanations
+            }
+
+        except Exception as e:
+            fallback = code + f"\n# [!] Enhancer crashed: {str(e)}"
+
+            return {
+                "enhanced_code": fallback,
+                "diff": create_diff(code, fallback),
+                "candidates": [],
+                "explanations": [{
+                    "change": "Error",
+                    "reason": str(e)
+                }]
+            }

requirements.txt

@@ -1,34 +1,28 @@
-
-# Flask dependencies
-flask==2.3.3
-flask-cors==4.0.0
-flask-jwt-extended==4.5.3
-flask-limiter==3.5.0
-flask-compress==1.13
-flask-caching==2.1.0
-python-dotenv==1.0.0
-
-# Database
-pymongo==4.5.0
-
-# Validation
-pydantic>=2.9.2
-
-
-# AI/ML dependencies
-torch>=2.0.0
-transformers>=4.30.0
-
-# Security scanning tools
-bandit>=1.7.5
-semgrep>=1.45.0
-
-# Optional: for better performance
-accelerate>=0.20.0
-safetensors>=0.3.0
-
-bcrypt>=4.0.1
-python-dotenv
-gunicorn
-
-requests==2.31.0
+flask==2.3.3
+flask-cors==4.0.0
+flask-jwt-extended==4.5.3
+flask-limiter==3.5.0
+flask-compress==1.13
+flask-caching==2.1.0
+python-dotenv==1.0.0
+
+pymongo==4.5.0
+
+pydantic>=2.9.2
+
+
+torch>=2.0.0
+transformers==4.40.2
+tokenizers==0.19.1
+
+bandit>=1.7.5
+semgrep>=1.45.0
+
+accelerate>=0.20.0
+safetensors>=0.3.0
+
+bcrypt>=4.0.1
+python-dotenv
+gunicorn
+
+requests==2.31.0

schemas.py

@@ -1,11 +1,43 @@
-# backend/schemas.py
-from pydantic import BaseModel
+from pydantic import BaseModel, field_validator
 from typing import List, Optional
 
 class FileModel(BaseModel):
     filename: str
     content: str
 
+    @field_validator('filename')
+    @classmethod
+    def validate_filename(cls, v):
+        if not v or not v.strip():
+            raise ValueError('Filename cannot be empty')
+        if len(v) > 255:
+            raise ValueError('Filename must be 255 characters or less')
+        return v
+
+    @field_validator('content')
+    @classmethod
+    def validate_content(cls, v):
+        if len(v) > 100_000:
+            raise ValueError('File content exceeds 100 KB limit. Please split into smaller files.')
+        return v
+
 class ScanRequest(BaseModel):
     files: List[FileModel]
     language: str = "python"
+
+    @field_validator('files')
+    @classmethod
+    def validate_files(cls, v):
+        if len(v) == 0:
+            raise ValueError('At least one file is required')
+        if len(v) > 10:
+            raise ValueError('Maximum 10 files allowed per scan request')
+        return v
+
+    @field_validator('language')
+    @classmethod
+    def validate_language(cls, v):
+        supported = ("python", "javascript")
+        if v.lower() not in supported:
+            raise ValueError(f'Language must be one of: {", ".join(supported)}')
+        return v.lower()