Upload folder using huggingface_hub

Dockerfile

@@ -1,14 +1,14 @@
-FROM python:3.9
-
-RUN useradd -m -u 1000 user
-USER user
-ENV PATH="/home/user/.local/bin:$PATH"
+FROM python:3.11-slim
 
 WORKDIR /app
 
-COPY --chown=user ./requirements.txt requirements.txt
-RUN pip install --no-cache-dir --upgrade -r requirements.txt
+RUN apt-get update && apt-get install -y gcc && rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY app.py .
 
-COPY --chown=user . /app
+EXPOSE 7860
 
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -1,11 +1,17 @@
 ---
-title: Polyguard Api
-emoji: 💻
-colorFrom: pink
-colorTo: gray
+title: PolyGuard API
+emoji: 🛡️
+colorFrom: blue
+colorTo: green
 sdk: docker
-pinned: false
-short_description: PolyGuard code security analyzer API
+app_port: 7860
+pinned: true
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# PolyGuard API
+
+Code vulnerability scanner powered by fine-tuned CodeBERT.
+
+- **Model**: MUHAMMADSAADAMIN/PolyGuard
+- **Docs**: /docs
+- **Analyze**: POST /analyze

app.py

@@ -1,99 +1,193 @@
-import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+import os, re, random
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
+from typing import List
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-MODEL_NAME = "MUHAMMADSAADAMIN/polyguard-model"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
+# ── Load model from HF Hub ─────────────────────────────────
+MODEL_ID  = "MUHAMMADSAADAMIN/PolyGuard"
+print(f"Loading model: {MODEL_ID}")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model     = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
 model.eval()
+print("✓ Model ready")
 
-suggestions = {
-    "sqli":    "Use parameterized queries instead of building SQL strings manually.",
-    "xss":     "Sanitize all user inputs before rendering them to the page.",
-    "secrets": "Never hardcode API keys or passwords. Use environment variables instead.",
-    "crypto":  "Avoid MD5 and SHA1. Use SHA256 or bcrypt for hashing.",
-    "memory":  "Always check buffer sizes before copying data in C/C++.",
-    "auth":    "Always verify user permissions before returning sensitive data.",
-}
-
-language_tips = {
+# ── Vulnerability rules ────────────────────────────────────
+VULN_RULES = {
     "python": [
-        "Use list comprehensions instead of for loops where possible.",
-        "Use f-strings for string formatting instead of .format() or %.",
-        "Use 'with open()' for file handling instead of open/close.",
+        (r"execute\s*\(\s*[f'\"].*?\{",          "Use parameterized queries instead of building SQL strings manually."),
+        (r"execute\s*\(\s*\".*?%",                  "Use parameterized queries instead of building SQL strings manually."),
+        (r"eval\s*\(",                                 "Avoid eval() — it executes arbitrary code and is a critical security risk."),
+        (r"exec\s*\(",                                 "Avoid exec() — it executes arbitrary code and is a critical security risk."),
+        (r"pickle\.loads?\s*\(",                      "Avoid pickle.load() on untrusted data — it can execute arbitrary code."),
+        (r"subprocess.*shell\s*=\s*True",              "Never use shell=True in subprocess — use a list of arguments instead."),
+        (r"os\.system\s*\(",                          "Avoid os.system() — use subprocess with a list of arguments instead."),
+        (r"hashlib\.md5\s*\(",                        "MD5 is cryptographically broken — use SHA-256 or stronger."),
+        (r"hashlib\.sha1\s*\(",                       "SHA-1 is weak — use SHA-256 or stronger."),
+        (r"random\.(random|randint|choice)\s*\(",     "Use secrets module instead of random for security-sensitive operations."),
+        (r"open\s*\(.*['\"]w['\"]",               "Validate file paths before writing to prevent path traversal attacks."),
+        (r"request\.(args|form|json)\[",               "Validate and sanitize all user input before use."),
+        (r"render_template_string\s*\(",               "Avoid render_template_string with user input — use template files instead."),
+        (r"yaml\.load\s*\(",                          "Use yaml.safe_load() instead of yaml.load() to prevent code execution."),
+        (r"SSL_VERIFY\s*=\s*False|verify\s*=\s*False", "Never disable SSL verification in production."),
+        (r"password\s*=\s*['\"][^'\"]{1,20}['\"]", "Hardcoded password detected — use environment variables instead."),
+        (r"secret\s*=\s*['\"][^'\"]{1,20}['\"]",   "Hardcoded secret detected — use environment variables instead."),
+        (r"api_key\s*=\s*['\"][^'\"]+['\"]",        "Hardcoded API key detected — use environment variables instead."),
     ],
     "javascript": [
-        "Use const and let instead of var.",
-        "Use async/await instead of nested callbacks.",
-        "Always use === instead of == for comparisons.",
+        (r"eval\s*\(",                                 "Avoid eval() — it executes arbitrary code and is a critical security risk."),
+        (r"innerHTML\s*=",                              "Avoid innerHTML — use textContent or DOMPurify to prevent XSS."),
+        (r"document\.write\s*\(",                     "Avoid document.write() — it can lead to XSS vulnerabilities."),
+        (r"dangerouslySetInnerHTML",                     "Avoid dangerouslySetInnerHTML — sanitize content with DOMPurify first."),
+        (r"localStorage\.setItem.*password",            "Never store passwords or secrets in localStorage."),
+        (r"Math\.random\s*\(",                        "Use crypto.getRandomValues() instead of Math.random() for security tokens."),
+        (r"http://",                                     "Use HTTPS instead of HTTP for all external requests."),
+        (r"password\s*=\s*['\"][^'\"]+['\"]",   "Hardcoded password detected — use environment variables instead."),
     ],
-    "java": [
-        "Use try-with-resources for handling streams and connections.",
-        "Use StringBuilder instead of String concatenation in loops.",
+    "sql": [
+        (r"'\s*\+\s*",                               "String concatenation in SQL is vulnerable to injection — use parameterized queries."),
+        (r"GRANT ALL",                                   "Avoid GRANT ALL — apply principle of least privilege."),
+        (r"DROP TABLE",                                  "Dangerous DDL statement detected — ensure proper access controls."),
+        (r"xp_cmdshell",                                 "xp_cmdshell is a critical security risk — disable it on the server."),
     ],
-    "go": [
-        "Always handle errors explicitly.",
-        "Use goroutines for concurrency instead of threads.",
+    "php": [
+        (r"mysql_query\s*\(",                          "mysql_* functions are deprecated — use PDO or mysqli with prepared statements."),
+        (r"\$_GET\[|\$_POST\[|\$_REQUEST\[",      "Sanitize all user input from $_GET/$_POST/$_REQUEST before use."),
+        (r"eval\s*\(",                                 "Avoid eval() — it executes arbitrary code."),
+        (r"system\s*\(|exec\s*\(",                   "Avoid system()/exec() with user input — use escapeshellarg()."),
+        (r"md5\s*\(",                                  "MD5 is not suitable for password hashing — use password_hash() instead."),
     ],
 }
 
-app = FastAPI(title="PolyGuard API")
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+CODE_TIPS = {
+    "python":     ["Use list comprehensions instead of for loops.", "Use f-strings for string formatting.", "Use with open() for file handling.", "Add type hints to function signatures.", "Use logging instead of print() in production.", "Use dataclasses or Pydantic instead of plain dicts."],
+    "javascript": ["Use const and let instead of var.", "Use async/await instead of callback chains.", "Use strict equality (===) instead of ==.", "Prefer arrow functions for concise syntax."],
+    "sql":        ["Always use parameterized queries.", "Add indexes on frequently queried columns.", "Use EXPLAIN to analyze query performance."],
+    "php":        ["Use Composer for dependency management.", "Enable strict_types=1 at the top of files.", "Use prepared statements for all database queries."],
+}
 
-class CodeRequest(BaseModel):
-    code:     str
-    language: str = "python"
+SMART_TIPS = {
+    "sql_injection":     "Use parameterized queries e.g. cursor.execute(query, params) to prevent SQL injection.",
+    "code_execution":    "Replace eval()/exec() with ast.literal_eval() for safe expression parsing.",
+    "hardcoded_secrets": "Store secrets in environment variables and load with os.environ.get(KEY).",
+    "weak_crypto":       "Replace MD5/SHA1 with hashlib.sha256() or bcrypt for password hashing.",
+    "xss":               "Sanitize output with DOMPurify or use textContent instead of innerHTML.",
+    "command_injection": "Pass arguments as a list to subprocess.run() and never use shell=True.",
+}
 
-@app.get("/")
-def home():
-    return {"status": "PolyGuard API is running!"}
+SUPPORTED_LANGUAGES = list(VULN_RULES.keys()) + ["java", "c", "cpp", "go", "ruby", "rust"]
 
-@app.post("/analyze")
-def analyze(request: CodeRequest):
-    inputs = tokenizer(
-        request.code,
-        return_tensors="pt",
-        truncation=True,
-        max_length=256,
-        padding=True
-    )
+def analyze_code(code: str, language: str):
+    language = language.lower().strip()
+    inputs   = tokenizer(code, return_tensors="pt", truncation=True, max_length=512, padding=True)
     with torch.no_grad():
-        outputs = model(**inputs)
+        logits = model(**inputs).logits
+
+    probs = torch.softmax(logits, dim=1).squeeze().tolist()
+    if not isinstance(probs, list):
+        probs = [probs, 1 - probs]
+
+    clean_conf = round(probs[0] * 100, 1)
+    vuln_conf  = round(probs[1] * 100, 1)
 
-    probs      = torch.softmax(outputs.logits, dim=1)
-    clean_conf = probs[0][0].item()
-    vuln_conf  = probs[0][1].item()
-    score      = round(clean_conf * 10, 1)
+    findings, matched_categories = [], set()
+    for pattern, message in VULN_RULES.get(language, []):
+        if re.search(pattern, code, re.IGNORECASE):
+            if message not in findings:
+                findings.append(message)
+                if "sql" in message.lower() or "query" in message.lower():
+                    matched_categories.add("sql_injection")
+                elif "eval" in message.lower() or "exec" in message.lower():
+                    matched_categories.add("code_execution")
+                elif any(w in message.lower() for w in ["password", "secret", "api_key"]):
+                    matched_categories.add("hardcoded_secrets")
+                elif any(w in message.lower() for w in ["md5", "sha1", "hash"]):
+                    matched_categories.add("weak_crypto")
+                elif "xss" in message.lower() or "innerhtml" in message.lower():
+                    matched_categories.add("xss")
+                elif "shell" in message.lower() or "subprocess" in message.lower():
+                    matched_categories.add("command_injection")
 
-    if score >= 8:
-        risk = "low"
-    elif score >= 5:
-        risk = "medium"
+    model_uncertain = abs(vuln_conf - clean_conf) < 15.0
+    if model_uncertain and len(findings) == 0:
+        score = 0.0
     else:
-        risk = "high"
+        finding_penalty = min(len(findings) * 0.8, 4.0)
+        base_score      = (vuln_conf / 100) * 6.0 if not model_uncertain else 0.0
+        score           = round(min(base_score + finding_penalty, 10.0), 1)
 
-    findings = []
-    if vuln_conf > 0.4:
-        findings.append(suggestions["sqli"])
-    if vuln_conf > 0.6:
-        findings.append(suggestions["xss"])
+    if score >= 7.0:   risk, verdict = "critical", "VULNERABLE"
+    elif score >= 4.5: risk, verdict = "high",     "VULNERABLE"
+    elif score >= 2.5: risk, verdict = "medium",   "REVIEW NEEDED"
+    elif score >= 1.0: risk, verdict = "low",      "MOSTLY SAFE"
+    else:              risk, verdict = "safe",      "CLEAN"
 
-    tips = language_tips.get(request.language.lower(), ["Keep learning!"])
+    tips    = [SMART_TIPS[c] for c in matched_categories if c in SMART_TIPS]
+    general = CODE_TIPS.get(language, ["Follow security best practices."])
+    tips   += random.sample(general, min(max(0, 3 - len(tips)), len(general)))
 
     return {
         "score":            score,
         "risk":             risk,
-        "verdict":          "CLEAN" if score >= 7 else "VULNERABLE",
-        "clean_confidence": round(clean_conf * 100, 1),
-        "vuln_confidence":  round(vuln_conf  * 100, 1),
+        "verdict":          verdict,
+        "clean_confidence": clean_conf,
+        "vuln_confidence":  vuln_conf,
         "findings":         findings,
-        "tips":             tips,
-    }
+        "tips":             tips[:3],
+        "language":         language,
+        "lines_analyzed":   len(code.splitlines()),
+    }
+
+# ── FastAPI app ────────────────────────────────────────────
+app = FastAPI(
+    title       = "PolyGuard API",
+    description = "Code vulnerability scanner powered by CodeBERT fine-tuned on CrossVUL.",
+    version     = "3.0.0",
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins     = ["*"],
+    allow_credentials = True,
+    allow_methods     = ["*"],
+    allow_headers     = ["*"],
+)
+
+class AnalyzeRequest(BaseModel):
+    code:     str
+    language: str
+
+class SnippetItem(BaseModel):
+    code:     str
+    language: str
+
+class BatchRequest(BaseModel):
+    snippets: List[SnippetItem]
+
+@app.get("/")
+def index():
+    return {
+        "service":             "PolyGuard — Code Vulnerability Scanner",
+        "version":             "3.0.0",
+        "model":               MODEL_ID,
+        "status":              "running",
+        "docs":                "/docs",
+        "supported_languages": SUPPORTED_LANGUAGES,
+        "endpoints":           ["/health", "/analyze", "/analyze_batch"],
+    }
+
+@app.get("/health")
+def health():
+    return {"status": "ok", "model": MODEL_ID, "supported_languages": SUPPORTED_LANGUAGES}
+
+@app.post("/analyze")
+def analyze(req: AnalyzeRequest):
+    return analyze_code(req.code, req.language)
+
+@app.post("/analyze_batch")
+def analyze_batch(req: BatchRequest):
+    results = [analyze_code(s.code, s.language) for s in req.snippets]
+    return {"count": len(results), "results": results}

requirements.txt

@@ -2,4 +2,4 @@ fastapi
 uvicorn
 transformers
 torch
-pydantic
+pydantic