diff --git a/.gitignore b/.gitignore index 8fc47ef..9c636b2 100644 --- a/.gitignore +++ b/.gitignore @@ -17,6 +17,7 @@ dist/ build/ .env .venv/ +.pytest_cache/ # Jupyter .ipynb_checkpoints/ diff --git a/README.md b/README.md index 45b3c01..c4304c3 100644 --- a/README.md +++ b/README.md @@ -401,13 +401,50 @@ pip install -r requirements.txt python src/training/train.py ``` +### Run the Demo App +```bash +# Needs a trained checkpoint at checkpoints/best_model.pt +python app.py + +# Or exercise the interface with untrained weights (no checkpoint required) +python app.py --smoke-test +``` + +Then open http://localhost:7860. `python app.py --help` lists the remaining flags +(`--checkpoint`, `--threshold`, `--host`, `--port`, `--share`). + +Model weights are not committed (`*.pt` is gitignored), so real inference needs a +trained or downloaded checkpoint. Checkpoints are validated on load: one that +does not match the architecture is a hard error rather than a silently partial +load. + +### Run Tests +```bash +pytest tests/ -v +``` + +The model tests build a tiny local encoder instead of downloading CodeBERT, so +the suite runs offline in about a second. + +### CWE Remediation Lookup +```bash +python -m src.models.fix_suggester --list +python -m src.models.fix_suggester CWE-119 89 79 +``` + ### Run Inference ```python +from src.models.loader import DEFAULT_CHECKPOINT_DIR, build_model, load_checkpoint from src.models.securescan_model import SecureScanModel from transformers import AutoTokenizer import torch -model = SecureScanModel() +# Built from the encoder config, then every parameter is taken from the +# checkpoint (load_checkpoint refuses anything below 100% coverage). +model = build_model(SecureScanModel, from_scratch=True) +load_checkpoint(model, DEFAULT_CHECKPOINT_DIR / 'best_model.pt') +model.eval() + tokenizer = AutoTokenizer.from_pretrained('microsoft/codebert-base') code = "char buf[10]; strcpy(buf, user_input);" inputs = tokenizer(code, return_tensors='pt', truncation=True, max_length=512) diff --git a/app.py b/app.py index 176119f..80be91b 100644 --- a/app.py +++ b/app.py @@ -1,92 +1,336 @@ """ -SecureScan AI Demo Application -Run locally: python app.py +SecureScan AI demo application. + +Run locally:: + + python app.py # uses checkpoints/best_model.pt + python app.py --checkpoint my.pt # explicit checkpoint + python app.py --smoke-test # UI only, untrained weights + Then open http://localhost:7860 + +Note on ``--smoke-test``: the model weights are deliberately random. It exists +so the interface can be exercised without training a model first, and the UI +labels the output as untrained so it is never mistaken for a real prediction. + +Checkpoints are validated by :mod:`src.models.loader` before use. One that does +not match the architecture is a hard error, because a partially applied +checkpoint leaves randomly-initialised layers that still answer confidently. """ +from __future__ import annotations + +import argparse import os -import torch -import gradio as gr -from transformers import AutoTokenizer, AutoModel import sys -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from pathlib import Path + +# Make ``src`` importable regardless of the working directory. +PROJECT_ROOT = Path(__file__).resolve().parent +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +import torch # noqa: E402 +from transformers import AutoTokenizer # noqa: E402 -from src.models.securescan_model import SecureScanModel +from src.models.loader import ( # noqa: E402 + DEFAULT_CHECKPOINT_DIR, + CheckpointError, + build_model, + load_checkpoint, +) +from src.models.securescan_model import SecureScanModel # noqa: E402 -# Load model and tokenizer -print("Loading model...") -tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base") -model = SecureScanModel() -model.load_state_dict(torch.load("src/models/baseline_mlp.pt", map_location="cpu"), strict=False) -model.eval() +# Defaults, overridable from the command line. +DEFAULT_CHECKPOINT = DEFAULT_CHECKPOINT_DIR / "best_model.pt" +TOKENIZER_NAME = "microsoft/codebert-base" +MAX_LENGTH = 512 +MAX_CHARS = 10_000 +# Decision threshold on P(vulnerable). Exposed because the default 0.5 is not +# necessarily right for a detector whose classes are heavily imbalanced. +DEFAULT_THRESHOLD = 0.5 -# Example code snippets -EXAMPLE_CODE = '''char buffer[10]; +EXAMPLE_CODE = """char buffer[10]; strcpy(buffer, input); -return 0;''' - -def predict_vulnerability(code: str) -> dict: - """Predict if code contains a vulnerability.""" - if not code or len(code.strip()) == 0: - return {"error": "Please provide valid code", "vulnerable": None, "confidence": None} - - if len(code) > 10000: - return {"error": "Code too long. Maximum 10,000 characters.", "vulnerable": None, "confidence": None} - - try: - inputs = tokenizer( - code, - return_tensors="pt", - truncation=True, - max_length=512 - ) - - with torch.no_grad(): - logits = model(inputs["input_ids"], inputs["attention_mask"]) - probs = torch.softmax(logits, dim=-1) - confidence = probs.max().item() - prediction = "Vulnerable" if logits.argmax().item() == 1 else "Safe" - - return { +return 0;""" + + +class Predictor: + """Wraps the tokenizer and model behind a single ``predict`` call. + + Args: + model: Trained (or, in smoke-test mode, untrained) classifier. + tokenizer: Matching tokenizer. + threshold: P(vulnerable) above which a snippet is flagged. + trained: False when the weights are known to be random. Surfaced in the + UI so an untrained model is never presented as a real result. + """ + + def __init__( + self, + model: torch.nn.Module, + tokenizer, + threshold: float = DEFAULT_THRESHOLD, + trained: bool = True, + ) -> None: + self.model = model.eval() + self.tokenizer = tokenizer + self.threshold = threshold + self.trained = trained + + def predict(self, code: str, threshold: float | None = None) -> dict: + """Classify a snippet. + + Args: + code: Source code to analyse. + threshold: Overrides the instance threshold for this call only. The + UI passes its slider value here rather than writing it onto the + predictor, which is shared across concurrent requests. + + Returns: + A JSON-serialisable result dict. ``error`` is non-None on failure. + """ + if not code or not code.strip(): + return _error("Please provide some source code.") + + if len(code) > MAX_CHARS: + return _error( + f"Code is {len(code):,} characters; the limit is " + f"{MAX_CHARS:,}. Analyse a single function instead." + ) + + active_threshold = self.threshold if threshold is None else float(threshold) + + try: + inputs = self.tokenizer( + code, + return_tensors="pt", + truncation=True, + max_length=MAX_LENGTH, + ) + with torch.no_grad(): + logits = self.model(inputs["input_ids"], inputs["attention_mask"]) + probs = torch.softmax(logits, dim=-1)[0] + # Read inside the guard: a model with an unexpected number of + # outputs should surface as an error, not as an IndexError trace. + p_vulnerable = float(probs[1]) + p_safe = float(probs[0]) + except Exception as exc: # surfaced to the UI, not a stack trace + return _error(f"Prediction failed: {exc}") + + # Truncation is measured in tokens, not whitespace-separated words: a + # single long token (a minified line, a long string literal) can exceed + # the window with very few spaces. A sequence that fills the window is + # reported as truncated, which errs towards warning the user. + truncated = int(inputs["input_ids"].shape[-1]) >= MAX_LENGTH + + flagged = p_vulnerable >= active_threshold + result = { "error": None, - "vulnerable": prediction, - "confidence": f"{confidence:.2%}" + "verdict": "Vulnerable" if flagged else "Safe", + "p_vulnerable": round(p_vulnerable, 4), + "p_safe": round(p_safe, 4), + "confidence": f"{max(p_vulnerable, 1 - p_vulnerable):.2%}", + "threshold": active_threshold, + "truncated": truncated, } - except Exception as e: - return {"error": f"Prediction failed: {str(e)}", "vulnerable": None, "confidence": None} - -# Create Gradio interface -with gr.Blocks(title="SecureScan AI") as demo: - gr.Markdown(""" - # SecureScan AI šŸ” - *AI-powered source code vulnerability detection* - - Trained on BigVul, DiverseVul, and FormAI datasets using CodeBERT + BiLSTM + MLP. - """) - - with gr.Row(): - with gr.Column(): - code_input = gr.Textbox( - label="Source Code", - placeholder="Paste C/C++ or Python code here...", - lines=10 + if not self.trained: + result["warning"] = ( + "UNTRAINED WEIGHTS ā€” this prediction is meaningless. " + "Rerun without --smoke-test using a trained checkpoint." ) - submit_btn = gr.Button("Analyze Vulnerability") - example_btn = gr.Button("Load Sample") - - with gr.Column(): - output = gr.JSON(label="Result") - - submit_btn.click(predict_vulnerability, inputs=code_input, outputs=output) - example_btn.click(lambda x: x, inputs=gr.Number(value=EXAMPLE_CODE, visible=False), outputs=code_input) - - gr.Markdown("---") - gr.Markdown(""" - ## About - - **Model**: CodeBERT + BiLSTM + MLP (F1: 0.9252) - - **Limitations**: May miss obfuscated vulnerabilities, best on C/C++ and Python - - **GitHub**: https://github.com/Salman1122334411/SecureScan-AI - """) + return result + + +def _error(message: str) -> dict: + """Build a result dict carrying only an error message.""" + return {"error": message, "verdict": None, "confidence": None} + + +def build_predictor( + checkpoint: str | os.PathLike[str] | None, + *, + smoke_test: bool = False, + threshold: float = DEFAULT_THRESHOLD, +) -> Predictor: + """Construct a :class:`Predictor` from a checkpoint or random weights. + + Args: + checkpoint: Checkpoint path. Ignored when ``smoke_test`` is True. + smoke_test: Build an untrained model instead of loading weights. + threshold: Decision threshold on P(vulnerable). + + Returns: + A ready-to-use predictor. + + Raises: + CheckpointError: If the checkpoint does not fit the architecture. + FileNotFoundError: If the checkpoint cannot be located. + """ + print(f"Loading tokenizer {TOKENIZER_NAME} ...") + tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME) + + # ``from_scratch`` builds the encoder from its published config instead of + # downloading ~500MB of pretrained weights that would be overwritten + # anyway: ``load_checkpoint`` refuses anything below 100% coverage, so every + # parameter ā€” encoder included ā€” is replaced by the checkpoint's. In + # smoke-test mode the random weights are the point. + model = build_model(SecureScanModel, from_scratch=True) + + if smoke_test: + print("SMOKE TEST MODE ā€” weights are random, predictions are meaningless.") + return Predictor(model, tokenizer, threshold, trained=False) + + path = checkpoint or DEFAULT_CHECKPOINT + print(f"Loading checkpoint {path} ...") + load_checkpoint(model, path) # raises if it does not fit + print("Checkpoint loaded.") + return Predictor(model, tokenizer, threshold, trained=True) + + +def build_interface(predictor: Predictor): + """Assemble the Gradio interface. + + Args: + predictor: Backend used by the Analyse button. + + Returns: + An unlaunched ``gr.Blocks`` instance. + """ + import gradio as gr + + banner = ( + "### :warning: Smoke-test mode\n" + "This instance is running **untrained, randomly-initialised weights**. " + "Predictions below are meaningless and are shown only to exercise the " + "interface." + if not predictor.trained + else "" + ) + + with gr.Blocks(title="SecureScan AI") as demo: + gr.Markdown( + """ + # SecureScan AI + *Source code vulnerability detection ā€” CodeBERT + BiLSTM + MLP* + + Paste a C/C++ or Python function to classify it as **Safe** or + **Vulnerable**. + """ + ) + if banner: + gr.Markdown(banner) + + with gr.Row(): + with gr.Column(): + code_input = gr.Textbox( + label="Source code", + placeholder="Paste a single function here...", + lines=12, + ) + with gr.Row(): + submit_btn = gr.Button("Analyse", variant="primary") + example_btn = gr.Button("Load sample") + threshold_slider = gr.Slider( + minimum=0.05, + maximum=0.95, + value=predictor.threshold, + step=0.05, + label="Decision threshold on P(vulnerable)", + info="Lower flags more code as vulnerable. Default 0.5.", + ) + with gr.Column(): + output = gr.JSON(label="Result") + + submit_btn.click( + fn=lambda code, thr: _predict_with(predictor, code, thr), + inputs=[code_input, threshold_slider], + outputs=output, + ) + # A no-input click handler. The previous version wired a string into a + # gr.Number input, which raised on every click. + example_btn.click(fn=lambda: EXAMPLE_CODE, inputs=None, outputs=code_input) + + gr.Markdown( + """ + --- + ### About + - **Model**: CodeBERT (6 of 12 layers frozen) + 2-layer BiLSTM + MLP + - **Scope**: binary classification only ā€” this build does not + predict a CWE category + - **Limitations**: best on C/C++ and Python; may miss obfuscated or + multi-function vulnerabilities; a single function is the intended + unit of analysis + - **Status**: see `docs/RESULTS.md` for the current, reproducible + evaluation. Metrics previously shown here were not reproducible + from this repository. + """ + ) + + return demo + + +def _predict_with(predictor: Predictor, code: str, threshold: float) -> dict: + """Apply a UI-selected threshold, then delegate to the predictor. + + The threshold travels as an argument rather than being written onto the + predictor, so two concurrent requests cannot overwrite each other's setting. + """ + return predictor.predict(code, threshold=threshold) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + """Parse command-line arguments.""" + parser = argparse.ArgumentParser(description="SecureScan AI demo app.") + parser.add_argument( + "--checkpoint", + default=None, + help=f"Path to a trained checkpoint (default: {DEFAULT_CHECKPOINT}).", + ) + parser.add_argument( + "--smoke-test", + action="store_true", + help="Run the UI with untrained weights to verify the interface.", + ) + parser.add_argument("--threshold", type=float, default=DEFAULT_THRESHOLD) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=7860) + parser.add_argument( + "--share", + action="store_true", + help="Create a public Gradio link. Off by default.", + ) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + """Entry point. Returns a process exit code.""" + args = parse_args(argv) + + try: + predictor = build_predictor( + args.checkpoint, + smoke_test=args.smoke_test, + threshold=args.threshold, + ) + except (CheckpointError, FileNotFoundError) as exc: + print(f"\nERROR: {exc}\n", file=sys.stderr) + print( + "No trained weights were found. Options:\n" + " 1. Train one: python -m src.training.train\n" + " 2. Download a release checkpoint into checkpoints/\n" + " 3. Verify the UI: python app.py --smoke-test\n", + file=sys.stderr, + ) + return 1 + + build_interface(predictor).launch( + server_name=args.host, + server_port=args.port, + share=args.share, + ) + return 0 + if __name__ == "__main__": - demo.launch(server_port=7860) \ No newline at end of file + raise SystemExit(main()) diff --git a/docs/DEPLOYMENT.md b/docs/DEPLOYMENT.md index 9471f9f..7721870 100644 --- a/docs/DEPLOYMENT.md +++ b/docs/DEPLOYMENT.md @@ -39,13 +39,14 @@ python src/training/train.py ### Using Pre-trained Model ```python -import torch +from src.models.loader import DEFAULT_CHECKPOINT_DIR, build_model, load_checkpoint from src.models.securescan_model import SecureScanModel from transformers import AutoTokenizer -# Load model -model = SecureScanModel() -model.load_state_dict(torch.load('src/models/baseline_mlp.pt')) +# Built from the encoder config here; every parameter is then taken from the +# checkpoint, which is validated before anything is applied. +model = build_model(SecureScanModel, from_scratch=True) +load_checkpoint(model, DEFAULT_CHECKPOINT_DIR / 'best_model.pt') model.eval() # Load tokenizer @@ -67,6 +68,12 @@ with torch.no_grad(): print(f"Prediction: {prediction}") ``` +Checkpoints are not committed (`*.pt` is gitignored): train one, or drop a +released checkpoint into `checkpoints/`. A checkpoint that does not fit the +architecture raises instead of leaving layers randomly initialised ā€” the snippet +that used to be here loaded an MLP checkpoint into this model, which silently +discarded every weight. + ### Command Line Interface ```bash @@ -77,11 +84,17 @@ python -c "from src.training.train import final_evaluation; ..." ## Testing ```bash -# Run unit tests +# Run unit tests (offline: a tiny local encoder stands in for CodeBERT) pytest tests/ -v -# Run specific test -python tests/test_pipeline.py +# Run a single module +pytest tests/test_loader.py -v +``` + +The demo app can be exercised without a checkpoint: + +```bash +python app.py --smoke-test # runs the UI with untrained weights ``` ## Project Structure for Deployment diff --git a/src/models/fix_suggester.py b/src/models/fix_suggester.py index cd65695..46cae77 100644 --- a/src/models/fix_suggester.py +++ b/src/models/fix_suggester.py @@ -1,62 +1,243 @@ +""" +CWE remediation lookup for SecureScan AI. +Maps a CWE identifier to a human-readable explanation and a concrete fix. + +The database path is resolved relative to this file rather than the current +working directory, so importing or running this module works from anywhere:: + + python -m src.models.fix_suggester CWE-119 + python -m src.models.fix_suggester 119 89 79 + python -m src.models.fix_suggester --list +""" + +from __future__ import annotations + +import argparse import json -import os +import sys +from pathlib import Path +from typing import Any, Dict, List + +# Resolve the database next to this module. The original implementation used a +# bare relative path, so importing from any other directory raised +# FileNotFoundError. +DATABASE_PATH = Path(__file__).resolve().parent / "cwe_database.json" + +SEVERITY_EMOJI = { + "Critical": "šŸ”“", + "High": "šŸŸ ", + "Medium": "šŸŸ”", + "Low": "šŸŸ¢", +} + +# Order used when rendering "list" output. +SEVERITY_ORDER = ["Critical", "High", "Medium", "Low"] + + +def load_cwe_database(path: str | Path = DATABASE_PATH) -> Dict[str, Dict[str, Any]]: + """Load the CWE remediation database. + + Args: + path: Location of the JSON database. + + Returns: + Mapping of CWE id to its remediation record. + + Raises: + FileNotFoundError: If the database is missing. + ValueError: If the file is not a JSON object. + """ + database_path = Path(path) + if not database_path.is_file(): + raise FileNotFoundError( + f"CWE database not found at {database_path}. " + "It ships with the repository at src/models/cwe_database.json." + ) + + with database_path.open(encoding="utf-8") as handle: + data = json.load(handle) + + if not isinstance(data, dict): + raise ValueError(f"{database_path} should contain a JSON object of CWE records.") + + return data + + +# Loaded once at import. A failure here is a genuine packaging error, so it is +# allowed to propagate rather than being swallowed into an empty dict. +CWE_DATABASE = load_cwe_database() + + +def normalise_cwe_id(cwe_id: str) -> str: + """Coerce a user-supplied identifier into ``CWE-`` form. + + Args: + cwe_id: e.g. ``"119"``, ``"cwe-119"``, ``" CWE-119 "``. + + Returns: + The canonical identifier, e.g. ``"CWE-119"``. + """ + cleaned = str(cwe_id).strip().upper() + if not cleaned.startswith("CWE-"): + cleaned = f"CWE-{cleaned}" + return cleaned + + +def get_fix_suggestion(cwe_id: str) -> Dict[str, Any]: + """Look up remediation guidance for a CWE identifier. + + Args: + cwe_id: CWE identifier in any of the accepted forms. -# Load CWE database -with open('cwe_database.json', 'r') as f: - CWE_DATABASE = json.load(f) + Returns: + A record with ``found=True`` and the database fields, or ``found=False`` + plus a ``message`` explaining the miss. + """ + canonical = normalise_cwe_id(cwe_id) -def get_fix_suggestion(cwe_id): - cwe_id = cwe_id.strip().upper() - if not cwe_id.startswith("CWE-"): - cwe_id = "CWE-" + cwe_id - if cwe_id not in CWE_DATABASE: + if canonical not in CWE_DATABASE: return { "found": False, - "cwe_id": cwe_id, - "message": f"CWE {cwe_id} not in database." + "cwe_id": canonical, + "message": ( + f"{canonical} is not in the database " + f"({len(CWE_DATABASE)} entries available; " + "run with --list to see them)." + ), } - data = CWE_DATABASE[cwe_id] + + data = CWE_DATABASE[canonical] return { "found": True, - "cwe_id": cwe_id, + "cwe_id": canonical, "name": data["name"], "severity": data["severity"], "description": data["description"], "why_dangerous": data["why_dangerous"], "bad_code": data["bad_code"], "good_code": data["good_code"], - "fix_steps": data["fix_steps"] + "fix_steps": data["fix_steps"], } -def print_fix_report(cwe_id): + +def format_fix_report(cwe_id: str) -> str: + """Render remediation guidance as a printable report. + + Args: + cwe_id: CWE identifier in any of the accepted forms. + + Returns: + The formatted report, or a short message if the CWE is unknown. + """ result = get_fix_suggestion(cwe_id) if not result["found"]: - print(f"Not found: {result['message']}") - return - severity_emoji = { - "Critical": "šŸ”“", - "High": "šŸŸ ", - "Medium": "šŸŸ”", - "Low": "šŸŸ¢" - } - emoji = severity_emoji.get(result["severity"], "āšŖ") - print("=" * 60) - print(f"šŸ” VULNERABILITY REPORT") - print("=" * 60) - print(f"CWE ID: {result['cwe_id']}") - print(f"Name: {result['name']}") - print(f"Severity: {emoji} {result['severity']}") - print("-" * 60) - print(f"\nšŸ“– WHAT IS IT?") - print(f"{result['description']}") - print(f"\nāš ļø WHY IS IT DANGEROUS?") - print(f"{result['why_dangerous']}") - print(f"\nāŒ VULNERABLE CODE:") - print(result['bad_code']) - print(f"\nāœ… FIXED CODE:") - print(result['good_code']) - print(f"\nšŸ”§ HOW TO FIX IT:") - for i, step in enumerate(result['fix_steps'], 1): - print(f" {i}. {step}") - print("=" * 60) + return f"Not found: {result['message']}" + + emoji = SEVERITY_EMOJI.get(result["severity"], "āšŖ") + lines = [ + "=" * 60, + "VULNERABILITY REPORT", + "=" * 60, + f"CWE ID: {result['cwe_id']}", + f"Name: {result['name']}", + f"Severity: {emoji} {result['severity']}", + "-" * 60, + "", + "WHAT IS IT?", + result["description"], + "", + "WHY IS IT DANGEROUS?", + result["why_dangerous"], + "", + "VULNERABLE CODE:", + result["bad_code"], + "", + "FIXED CODE:", + result["good_code"], + "", + "HOW TO FIX IT:", + ] + lines += [f" {i}. {step}" for i, step in enumerate(result["fix_steps"], 1)] + lines.append("=" * 60) + return "\n".join(lines) + + +def print_fix_report(cwe_id: str) -> None: + """Print :func:`format_fix_report` to stdout. + + Args: + cwe_id: CWE identifier in any of the accepted forms. + """ + print(format_fix_report(cwe_id)) + + +def list_cwes() -> List[Dict[str, str]]: + """Summarise the database, ordered by severity then identifier. + + Returns: + One ``{"cwe_id", "name", "severity"}`` record per entry. + """ + order = {name: i for i, name in enumerate(SEVERITY_ORDER)} + records = [ + { + "cwe_id": cwe_id, + "name": record["name"], + "severity": record["severity"], + } + for cwe_id, record in CWE_DATABASE.items() + ] + return sorted( + records, + key=lambda r: (order.get(r["severity"], len(order)), r["cwe_id"]), + ) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + """Parse command-line arguments.""" + parser = argparse.ArgumentParser( + prog="python -m src.models.fix_suggester", + description="Look up CWE remediation guidance.", + ) + parser.add_argument( + "cwe_ids", + nargs="*", + help="One or more CWE identifiers, e.g. CWE-119 or 119.", + ) + parser.add_argument( + "--list", + action="store_true", + help="List every CWE in the database and exit.", + ) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + """Entry point. Returns a process exit code.""" + args = parse_args(argv) + + if args.list: + for record in list_cwes(): + emoji = SEVERITY_EMOJI.get(record["severity"], "āšŖ") + print(f"{record['cwe_id']:<10} {emoji} {record['severity']:<9} {record['name']}") + return 0 + + if not args.cwe_ids: + print( + "No CWE id given. Try: python -m src.models.fix_suggester --list", + file=sys.stderr, + ) + return 2 + + exit_code = 0 + for index, cwe_id in enumerate(args.cwe_ids): + if index: + print() + print(format_fix_report(cwe_id)) + if not get_fix_suggestion(cwe_id)["found"]: + exit_code = 1 + return exit_code + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/models/loader.py b/src/models/loader.py new file mode 100644 index 0000000..adee9d6 --- /dev/null +++ b/src/models/loader.py @@ -0,0 +1,330 @@ +""" +Checkpoint loading utilities for SecureScan AI. + +This module exists because the original demo loaded a ``BaselineMLP`` +checkpoint into a ``SecureScanModel`` with ``strict=False``. None of the +parameter names overlap between those two architectures, so *every* weight was +silently discarded and the classifier head ran on random initialisation while +still reporting a confident probability. + +Loading is therefore treated as a validated operation: a checkpoint that does +not cover the model's parameters is an error, not a warning. +""" + +from __future__ import annotations + +import os +import warnings +from pathlib import Path +from typing import Any, Dict, Mapping, MutableMapping, Tuple, Type + +import torch +import torch.nn as nn + +# Repository root, resolved from this file rather than the current working +# directory so that imports and relative asset paths behave identically no +# matter where a script is invoked from. +PROJECT_ROOT = Path(__file__).resolve().parents[2] + +# Canonical on-disk locations. ``*.pt`` is gitignored, so these files are +# produced by training or fetched from a GitHub release rather than committed. +DEFAULT_CHECKPOINT_DIR = PROJECT_ROOT / "checkpoints" +LEGACY_CHECKPOINT_DIR = PROJECT_ROOT / "src" / "models" + + +class CheckpointError(RuntimeError): + """Raised when a checkpoint cannot be loaded into a model safely.""" + + +def resolve_checkpoint(path: str | os.PathLike[str]) -> Path: + """Resolve a checkpoint path, tolerating the legacy ``src/models`` layout. + + Args: + path: Absolute path, or a path relative to the project root. + + Returns: + The resolved absolute path. + + Raises: + FileNotFoundError: If no candidate location holds the checkpoint. + """ + candidate = Path(path) + if candidate.is_absolute() and candidate.is_file(): + return candidate + + searched = [] + for base in (PROJECT_ROOT, DEFAULT_CHECKPOINT_DIR, LEGACY_CHECKPOINT_DIR): + resolved = (base / candidate).resolve() + if resolved.is_file(): + return resolved + searched.append(resolved) + + raise FileNotFoundError( + f"Checkpoint {str(path)!r} not found. Searched:\n" + + "\n".join(f" - {p}" for p in searched) + + "\n\nModel weights are not committed to git (*.pt is gitignored). " + "Train a model first (see README 'Run Training'), or download a " + "released checkpoint into checkpoints/." + ) + + +def extract_state_dict(checkpoint: Any) -> Mapping[str, torch.Tensor]: + """Pull a ``state_dict`` out of a raw or wrapped checkpoint. + + Handles the three shapes produced across this project's history: a bare + ``state_dict``, the ``{"model_state_dict": ...}`` wrapper written by + :func:`src.utils.helpers.save_checkpoint`, and ``{"state_dict": ...}``. + + Args: + checkpoint: Object loaded from a ``.pt`` file. + + Returns: + A mapping of parameter name to tensor. + + Raises: + CheckpointError: If no recognisable state dict is present. + """ + if isinstance(checkpoint, nn.Module): + return checkpoint.state_dict() + + if isinstance(checkpoint, Mapping): + for key in ("model_state_dict", "state_dict", "model"): + inner = checkpoint.get(key) + if isinstance(inner, nn.Module): + return inner.state_dict() + if isinstance(inner, Mapping): + return inner + # A bare state dict: every value is a tensor. + if checkpoint and all(isinstance(v, torch.Tensor) for v in checkpoint.values()): + return checkpoint + + raise CheckpointError( + "Unrecognised checkpoint format. Expected a state_dict, or a mapping " + "containing one of 'model_state_dict', 'state_dict', or 'model'." + ) + + +def _normalise(module: nn.Module) -> nn.Module: + """Strip a ``DataParallel``/``DistributedDataParallel`` ``module.`` prefix.""" + return module.module if hasattr(module, "module") else module + + +def diagnose_state_dict( + model: nn.Module, + state_dict: Mapping[str, torch.Tensor], +) -> Tuple[Dict[str, torch.Tensor], list[str], list[str], list[str]]: + """Compare a checkpoint against a model's parameters. + + Args: + model: The model the checkpoint is intended for. + state_dict: Candidate parameters, possibly carrying a ``module.`` prefix. + + Returns: + Tuple of ``(usable, missing, unexpected, shape_mismatch)`` where + ``usable`` holds only the entries that match by name *and* shape, and + the three lists hold human-readable names of everything rejected. + """ + target = _normalise(model) + own = target.state_dict() + + # Normalise the incoming keys once, so a DataParallel checkpoint still maps. + incoming = { + (k[len("module.") :] if k.startswith("module.") else k): v + for k, v in state_dict.items() + } + + usable: Dict[str, torch.Tensor] = {} + missing: list[str] = [] + shape_mismatch: list[str] = [] + + for name, tensor in own.items(): + if name not in incoming: + missing.append(name) + continue + candidate = incoming[name] + if not isinstance(candidate, torch.Tensor): + shape_mismatch.append(f"{name} (not a tensor)") + continue + if tuple(candidate.shape) != tuple(tensor.shape): + shape_mismatch.append( + f"{name} (checkpoint {tuple(candidate.shape)} " + f"!= model {tuple(tensor.shape)})" + ) + continue + usable[name] = candidate + + unexpected = [name for name in incoming if name not in own] + + return usable, missing, unexpected, shape_mismatch + + +def load_checkpoint( + model: nn.Module, + path: str | os.PathLike[str], + *, + strict: bool = True, + min_coverage: float | None = None, + map_location: str | torch.device = "cpu", +) -> nn.Module: + """Load a checkpoint into ``model``, validating that it actually applies. + + Args: + model: Model to load weights into, modified in place. + path: Checkpoint location; see :func:`resolve_checkpoint`. + strict: If True, require full coverage and no shape mismatches. + min_coverage: Minimum fraction of model parameters the checkpoint must + supply. Defaults to 1.0 when ``strict`` is True and 0.0 otherwise, + so a deliberately partial load does not have to restate the + threshold. An explicit value is always honoured. + map_location: Device mapping passed to :func:`torch.load`. + + Returns: + The same ``model`` instance, for convenient chaining. + + Raises: + ValueError: If ``min_coverage`` is outside ``[0, 1]``. + FileNotFoundError: If the checkpoint cannot be located. + CheckpointError: If the checkpoint is missing parameters, contains + mismatched shapes, or falls below ``min_coverage``. The message + lists the offending keys so the mismatch is diagnosable. + + Warns: + RuntimeWarning: When a partial load was explicitly permitted and some + parameters are therefore left at their initial values. + """ + if min_coverage is None: + min_coverage = 1.0 if strict else 0.0 + if not 0.0 <= min_coverage <= 1.0: + raise ValueError( + f"min_coverage={min_coverage} must be within [0, 1]. Use 1.0 to " + "require a complete load, or strict=False to permit a partial one." + ) + + resolved = resolve_checkpoint(path) + + # weights_only=True prevents arbitrary code execution from a pickle. This + # is also the default from PyTorch 2.6 onward, where omitting it breaks. + checkpoint = torch.load(resolved, map_location=map_location, weights_only=True) + state_dict = extract_state_dict(checkpoint) + + usable, missing, unexpected, shape_mismatch = diagnose_state_dict(model, state_dict) + + total = len(_normalise(model).state_dict()) + coverage = len(usable) / total if total else 0.0 + + if strict and (missing or shape_mismatch): + details = [] + if missing: + details.append( + f"{len(missing)} parameter(s) absent from the checkpoint, e.g. " + + ", ".join(missing[:5]) + ) + if shape_mismatch: + details.append( + f"{len(shape_mismatch)} shape mismatch(es), e.g. " + + ", ".join(shape_mismatch[:3]) + ) + raise CheckpointError( + f"Checkpoint {resolved.name!r} does not fit the model " + f"({coverage:.1%} coverage).\n " + + "\n ".join(details) + + "\n\nThis usually means the checkpoint was saved from a different " + "architecture (for example a BaselineMLP checkpoint loaded into " + "SecureScanModel). Refusing to load, because a partial load leaves " + "randomly-initialised layers that still produce confident output." + ) + + if coverage < min_coverage: + raise CheckpointError( + f"Checkpoint {resolved.name!r} covers only {coverage:.1%} of model " + f"parameters, below the required {min_coverage:.1%}." + ) + + if missing or shape_mismatch: + warnings.warn( + f"Partial load of {resolved.name!r}: {len(missing)} parameter(s) " + f"and {len(shape_mismatch)} shape mismatch(es) were left at their " + "initial values, so output from this model is not trustworthy.", + RuntimeWarning, + stacklevel=2, + ) + + _normalise(model).load_state_dict(usable, strict=False) + return model + + +def save_checkpoint( + model: nn.Module, + path: str | os.PathLike[str], + *, + epoch: int | None = None, + metrics: Mapping[str, float] | None = None, + optimizer: torch.optim.Optimizer | None = None, + extra: Mapping[str, Any] | None = None, +) -> Path: + """Save a checkpoint in the wrapper format this project reads back. + + Args: + model: Model whose weights are written. + path: Destination path; parent directories are created. + epoch: Optional epoch number for provenance. + metrics: Optional validation metrics for provenance. + optimizer: Optional optimizer state, restored by the training loop. + extra: Optional additional provenance (config, git SHA, dataset hash). + + Returns: + The path written to. + """ + destination = Path(path) + if not destination.is_absolute(): + destination = (PROJECT_ROOT / destination).resolve() + destination.parent.mkdir(parents=True, exist_ok=True) + + payload: MutableMapping[str, Any] = { + "model_state_dict": _normalise(model).state_dict(), + } + if epoch is not None: + payload["epoch"] = epoch + if metrics is not None: + payload["metrics"] = dict(metrics) + if optimizer is not None: + payload["optimizer_state_dict"] = optimizer.state_dict() + if extra is not None: + payload.update(extra) + + torch.save(payload, destination) + return destination + + +def build_model( + model_cls: Type[nn.Module], + *, + aliases: Mapping[str, str] | None = None, + **kwargs: Any, +) -> nn.Module: + """Instantiate ``model_cls``, forwarding only the kwargs it accepts. + + Lets callers pass a full config dict without each model class having to + declare every field. + + Args: + model_cls: Model class to instantiate. + aliases: Optional ``{config_key: constructor_arg}`` renaming applied + before the constructor signature is consulted. Without it a config + key that only *nearly* matches is dropped silently and the default + is used instead ā€” ``config.yaml`` in this repository spells the + encoder ``codebert`` while the argument is ``codebert_model``. + **kwargs: Candidate constructor arguments; unknown keys are ignored. + + Returns: + An instance of ``model_cls``. + """ + import inspect + + renamed = dict(kwargs) + for source, target in (aliases or {}).items(): + if source in renamed and target not in renamed: + renamed[target] = renamed.pop(source) + + accepted = set(inspect.signature(model_cls.__init__).parameters) - {"self"} + return model_cls(**{k: v for k, v in renamed.items() if k in accepted}) diff --git a/src/models/securescan_model.py b/src/models/securescan_model.py index 34c7dbf..664f24f 100644 --- a/src/models/securescan_model.py +++ b/src/models/securescan_model.py @@ -1,28 +1,40 @@ """ Main SecureScan AI model: CodeBERT + BiLSTM + MLP. + Primary vulnerability detection architecture. """ +from __future__ import annotations + +from typing import Any, Dict, cast + import torch import torch.nn as nn -from transformers import AutoModel +from transformers import AutoConfig, AutoModel class SecureScanModel(nn.Module): """CodeBERT + Bidirectional LSTM + MLP for vulnerability detection. - + Architecture: 1. CodeBERT: Pretrained transformer for code understanding 2. BiLSTM: Captures sequential patterns in code 3. MLP: Final classification head - + Args: - codebert_model: Pretrained model name. - freeze_layers: Number of CodeBERT layers to freeze. + codebert_model: Pretrained model name, or a path to a local directory. + freeze_layers: Number of leading CodeBERT layers to freeze. lstm_hidden: LSTM hidden size per direction. lstm_layers: Number of LSTM layers. dropout: Dropout probability. num_classes: Number of output classes. + from_scratch: Build a randomly-initialised encoder from the model's + published config instead of downloading its weights. Intended for + tests and CI, where pulling ~500MB per run is not acceptable. The + resulting model is architecturally identical but untrained. + + Raises: + ValueError: If ``freeze_layers`` exceeds the encoder's layer count. """ def __init__( @@ -32,27 +44,42 @@ def __init__( lstm_hidden: int = 256, lstm_layers: int = 2, dropout: float = 0.3, - num_classes: int = 2 + num_classes: int = 2, + from_scratch: bool = False, ) -> None: - super(SecureScanModel, self).__init__() + super().__init__() + + # Encoder. ``from_scratch`` reuses the published config but skips the + # weight download, which keeps the parameter shapes identical. + if from_scratch: + config = AutoConfig.from_pretrained(codebert_model) + self.codebert = AutoModel.from_config(config) + else: + self.codebert = AutoModel.from_pretrained(codebert_model) - # CodeBERT - self.codebert = AutoModel.from_pretrained(codebert_model) + layer_count = len(self.codebert.encoder.layer) + if not 0 <= freeze_layers <= layer_count: + raise ValueError( + f"freeze_layers={freeze_layers} is out of range for an encoder " + f"with {layer_count} layers." + ) for i in range(freeze_layers): for param in self.codebert.encoder.layer[i].parameters(): param.requires_grad = False - # BiLSTM + # Read the encoder width from its config rather than hard-coding 768, + # so the model still builds against a different checkpoint. + encoder_dim = self.codebert.config.hidden_size + self.lstm = nn.LSTM( - input_size=768, + input_size=encoder_dim, hidden_size=lstm_hidden, num_layers=lstm_layers, batch_first=True, bidirectional=True, - dropout=dropout + dropout=dropout if lstm_layers > 1 else 0.0, ) - # MLP Classifier self.classifier = nn.Sequential( nn.Linear(lstm_hidden * 2, 256), nn.ReLU(), @@ -61,26 +88,56 @@ def __init__( nn.Linear(256, 128), nn.ReLU(), nn.Dropout(dropout), - nn.Linear(128, num_classes) + nn.Linear(128, num_classes), ) def forward( self, input_ids: torch.Tensor, - attention_mask: torch.Tensor + attention_mask: torch.Tensor, ) -> torch.Tensor: """Forward pass through the network. - + Args: - input_ids: Tokenized input IDs. - attention_mask: Attention mask tensor. - + input_ids: Tokenized input IDs, shape ``(batch, seq_len)``. + attention_mask: Attention mask tensor, same shape. + Returns: - Logits tensor of shape (batch_size, num_classes). + Logits of shape ``(batch, num_classes)``. """ codebert_out = self.codebert( input_ids=input_ids, - attention_mask=attention_mask + attention_mask=attention_mask, ) lstm_out, _ = self.lstm(codebert_out.last_hidden_state) + # The first position carries the [CLS] representation, which the + # frozen-layer scheme in this project was trained against. Note that + # for a bidirectional LSTM this is *not* a pooled summary of the + # sequence; changing it would invalidate existing checkpoints. return self.classifier(lstm_out[:, 0, :]) + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "SecureScanModel": + """Build a model from a config mapping, ignoring unknown keys. + + Args: + config: Mapping of constructor arguments, e.g. the ``model`` section + of ``phases/phase4-refinement/config.yaml``. + + Returns: + An initialised model. + + Note: + ``codebert`` is accepted as an alias for ``codebert_model``, because + that is how this project's config files spell it. Without the alias + the key would be dropped and the encoder would silently fall back to + the default. + """ + # Imported here rather than at module scope so this module stays + # importable on its own, and so no import cycle can form later. + from src.models.loader import build_model + + return cast( + "SecureScanModel", + build_model(cls, aliases={"codebert": "codebert_model"}, **config), + ) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..d33e2cc --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,81 @@ +"""Shared test fixtures and import setup. + +Two things are handled here. + +**Import path.** The repository root is added to ``sys.path`` so that both +``from src.models... import ...`` and ``import app`` work no matter how pytest +was invoked (``pytest tests/`` does not guarantee the root directory is +importable). + +**Offline models.** ``SecureScanModel()`` with its default +``codebert_model="microsoft/codebert-base"`` downloads roughly 500MB. Tests must +not do that, so the model fixtures point at a tiny local BERT config instead and +build with ``from_scratch=True``. Same code path, no network, milliseconds. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from transformers import BertConfig # noqa: E402 (needs sys.path set up first) + + +# Small enough to build and run instantly, but structurally a real BERT: every +# attribute SecureScanModel reaches for (``encoder.layer``, ``config.hidden_size``) +# exists, so the architecture under test is the production one. +TINY_VOCAB_SIZE = 64 +TINY_HIDDEN_SIZE = 32 +TINY_LAYERS = 2 + + +@pytest.fixture(scope="session") +def tiny_encoder_dir(tmp_path_factory) -> Path: + """Directory holding a minimal encoder config. + + ``AutoConfig.from_pretrained`` reads this from disk, so no network access is + needed; ``from_scratch=True`` then skips the weight download entirely. + """ + directory = tmp_path_factory.mktemp("tiny-codebert") + config = BertConfig( + vocab_size=TINY_VOCAB_SIZE, + hidden_size=TINY_HIDDEN_SIZE, + num_hidden_layers=TINY_LAYERS, + num_attention_heads=4, + intermediate_size=37, + max_position_embeddings=64, + ) + config.save_pretrained(directory) + return directory + + +@pytest.fixture +def make_tiny_model(tiny_encoder_dir): + """Factory returning fresh tiny models, for load-into-a-new-instance tests.""" + from src.models.securescan_model import SecureScanModel + + def _make(**overrides): + kwargs = { + "codebert_model": str(tiny_encoder_dir), + "freeze_layers": 0, + "lstm_hidden": 8, + "lstm_layers": 1, + "dropout": 0.0, + "from_scratch": True, + } + kwargs.update(overrides) + return SecureScanModel(**kwargs) + + return _make + + +@pytest.fixture +def tiny_model(make_tiny_model): + """An untrained SecureScanModel backed by the tiny encoder.""" + return make_tiny_model() \ No newline at end of file diff --git a/tests/test_app.py b/tests/test_app.py new file mode 100644 index 0000000..d6c6f5d --- /dev/null +++ b/tests/test_app.py @@ -0,0 +1,222 @@ +"""Tests for the Gradio demo in ``app.py``. + +Importing ``app`` must stay free of side effects: the previous version built the +tokenizer and the model at import time, which downloaded ~500MB and crashed when +no checkpoint was present. + +Prediction is exercised through a stub tokenizer and stub models, so these tests +need no network access and no trained checkpoint. +""" + +from __future__ import annotations + +import torch +import pytest + +import app +from app import Predictor, parse_args + + +class StubTokenizer: + """Stand-in for ``AutoTokenizer``, returning a deterministic encoding. + + ``length`` models the real tokenizer's truncation behaviour: it is capped at + ``max_length`` whenever a cap is supplied. + """ + + def __init__(self, length: int = 8) -> None: + self.length = length + self.calls: list[str] = [] + + def __call__(self, code, return_tensors=None, truncation=None, max_length=None): + self.calls.append(code) + length = self.length if max_length is None else min(self.length, max_length) + return { + "input_ids": torch.zeros(1, length, dtype=torch.long), + "attention_mask": torch.ones(1, length, dtype=torch.long), + } + + +class ConstantLogits(torch.nn.Module): + """Model stub returning the same logits for every input.""" + + def __init__(self, safe: float = 0.0, vulnerable: float = 0.0) -> None: + super().__init__() + self.values = (safe, vulnerable) + + def forward(self, input_ids, attention_mask=None): + base = torch.tensor(self.values, dtype=torch.float32) + return base.repeat(input_ids.shape[0], 1) + + +# softmax([2, -2]) puts P(vulnerable) at ~0.018, well inside "Safe". +SAFE_LOGITS = {"safe": 2.0, "vulnerable": -2.0} +VULNERABLE_LOGITS = {"safe": -2.0, "vulnerable": 2.0} + + +def make_predictor(logits=None, **kwargs) -> Predictor: + """Build a predictor around a stub model: no checkpoint, no network.""" + logits = SAFE_LOGITS if logits is None else logits + return Predictor(ConstantLogits(**logits), StubTokenizer(), **kwargs) + + +def test_importing_app_builds_nothing(): + """The module must not construct a model or tokenizer at import time.""" + assert not hasattr(app, "model") + assert not hasattr(app, "tokenizer") + + +def test_rejects_empty_input(): + predictor = make_predictor() + result = predictor.predict(" \n\t ") + assert result["error"] == "Please provide some source code." + assert result["verdict"] is None + assert predictor.tokenizer.calls == [] + + +def test_rejects_over_long_input_without_tokenizing(): + predictor = make_predictor() + result = predictor.predict("x" * (app.MAX_CHARS + 1)) + assert "limit is" in result["error"] + assert predictor.tokenizer.calls == [] + + +def test_clean_snippet_is_reported_safe(): + result = make_predictor().predict("int main() { return 0; }") + assert result["error"] is None + assert result["verdict"] == "Safe" + assert result["p_vulnerable"] < 0.5 + assert result["truncated"] is False + assert result["confidence"].endswith("%") + assert "warning" not in result + + +def test_flagged_snippet_is_reported_vulnerable(): + result = make_predictor(VULNERABLE_LOGITS).predict("strcpy(buffer, input);") + assert result["verdict"] == "Vulnerable" + assert result["p_vulnerable"] > 0.5 + + +def test_threshold_override_does_not_mutate_the_predictor(): + """A per-call threshold replaces shared-state mutation, which is unsafe when + Gradio serves concurrent requests.""" + predictor = make_predictor(SAFE_LOGITS, threshold=0.5) + + flagged = predictor.predict("int x;", threshold=0.01) + assert flagged["verdict"] == "Vulnerable" + assert flagged["threshold"] == 0.01 + + assert predictor.threshold == 0.5 + assert predictor.predict("int x;")["verdict"] == "Safe" + + +def test_predict_with_forwards_the_ui_threshold(): + """The Gradio handler must not write onto the shared predictor.""" + predictor = make_predictor(SAFE_LOGITS, threshold=0.5) + assert app._predict_with(predictor, "int x;", 0.01)["verdict"] == "Vulnerable" + assert predictor.threshold == 0.5 + + +def test_untrained_predictor_labels_its_output(): + result = make_predictor(trained=False).predict("int x;") + assert result["error"] is None + assert "UNTRAINED" in result["warning"] + + +def test_truncation_is_measured_on_tokens(): + """A long unbroken line must trip the flag: whitespace counting missed it.""" + at_cap = Predictor(ConstantLogits(), StubTokenizer(length=app.MAX_LENGTH)) + assert at_cap.predict("x" * 5000)["truncated"] is True + + below_cap = Predictor(ConstantLogits(), StubTokenizer(length=app.MAX_LENGTH - 1)) + assert below_cap.predict("int x;")["truncated"] is False + + +def test_model_failure_is_returned_as_an_error(): + class Broken(torch.nn.Module): + def forward(self, input_ids, attention_mask=None): + raise RuntimeError("boom") + + result = Predictor(Broken(), StubTokenizer()).predict("int x;") + assert "Prediction failed" in result["error"] + assert result["verdict"] is None + + +def test_unexpected_output_width_is_returned_as_an_error(): + """Indexing a single logit must not escape as an IndexError trace.""" + + class OneClass(torch.nn.Module): + def forward(self, input_ids, attention_mask=None): + return torch.zeros(input_ids.shape[0], 1) + + result = Predictor(OneClass(), StubTokenizer()).predict("int x;") + assert "Prediction failed" in result["error"] + assert result["verdict"] is None + + +def test_smoke_test_mode_skips_weights_and_files(monkeypatch, tiny_model): + """``--smoke-test`` must not download weights or read a checkpoint.""" + captured = {} + + def fake_build(model_cls, **kwargs): + captured.update(kwargs) + return tiny_model + + def unexpected(*args, **kwargs): + raise AssertionError("load_checkpoint must not run in smoke-test mode") + + monkeypatch.setattr(app, "build_model", fake_build) + monkeypatch.setattr(app, "load_checkpoint", unexpected) + monkeypatch.setattr( + app.AutoTokenizer, "from_pretrained", lambda *a, **k: StubTokenizer() + ) + + predictor = app.build_predictor(None, smoke_test=True) + + # The encoder is built from its config, so no ~500MB weight download happens. + assert captured == {"from_scratch": True} + assert predictor.trained is False + + +def test_checkpoint_mode_loads_the_requested_path(monkeypatch, tiny_model): + loaded = [] + + monkeypatch.setattr(app, "build_model", lambda *a, **k: tiny_model) + monkeypatch.setattr( + app, "load_checkpoint", lambda model, path, **kwargs: loaded.append(str(path)) + ) + monkeypatch.setattr( + app.AutoTokenizer, "from_pretrained", lambda *a, **k: StubTokenizer() + ) + + checkpoint = "my/own/checkpoint.pt" + predictor = app.build_predictor(checkpoint) + + assert loaded == [checkpoint] + assert predictor.trained is True + + +def test_parse_args_defaults(): + args = parse_args([]) + assert args.checkpoint is None + assert args.smoke_test is False + assert args.threshold == app.DEFAULT_THRESHOLD + assert (args.host, args.port, args.share) == ("127.0.0.1", 7860, False) + + +def test_parse_args_overrides(): + args = parse_args( + [ + "--checkpoint", + "cp.pt", + "--smoke-test", + "--threshold", + "0.7", + "--port", + "9000", + ] + ) + assert args.checkpoint == "cp.pt" + assert args.smoke_test is True + assert args.threshold == pytest.approx(0.7) + assert args.port == 9000 \ No newline at end of file diff --git a/tests/test_fix_suggester.py b/tests/test_fix_suggester.py new file mode 100644 index 0000000..04b2d2e --- /dev/null +++ b/tests/test_fix_suggester.py @@ -0,0 +1,116 @@ +"""Tests for the CWE remediation lookup in :mod:`src.models.fix_suggester`. + +The module used to open ``cwe_database.json`` through a bare relative path, so +importing it from anywhere but ``src/models`` raised ``FileNotFoundError``. +``test_database_path_is_independent_of_cwd`` covers that. +""" + +from __future__ import annotations + +import importlib +import sys + +import pytest + +from src.models import fix_suggester +from src.models.fix_suggester import ( + CWE_DATABASE, + format_fix_report, + get_fix_suggestion, + list_cwes, + main, + normalise_cwe_id, +) + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ("119", "CWE-119"), + ("cwe-119", "CWE-119"), + ("CWE-119", "CWE-119"), + (" Cwe-119 ", "CWE-119"), + ], +) +def test_normalise_cwe_id(raw, expected): + assert normalise_cwe_id(raw) == expected + + +def test_database_is_loaded(): + """A missing or empty file would silently turn every lookup into a miss.""" + assert CWE_DATABASE + assert fix_suggester.DATABASE_PATH.is_file() + + +def test_known_cwe_returns_remediation(): + result = get_fix_suggestion("119") + assert result["found"] is True + assert result["cwe_id"] == "CWE-119" + for field in ( + "name", + "severity", + "description", + "why_dangerous", + "bad_code", + "good_code", + "fix_steps", + ): + assert result[field], field + + +def test_unknown_cwe_reports_a_miss(): + result = get_fix_suggestion("CWE-999999") + assert result["found"] is False + assert "not in the database" in result["message"] + + +def test_list_cwes_is_sorted_by_severity(): + records = list_cwes() + assert records + + order = {name: i for i, name in enumerate(fix_suggester.SEVERITY_ORDER)} + severities = [order[record["severity"]] for record in records] + assert severities == sorted(severities) + + +def test_format_report_contains_every_section(): + report = format_fix_report("119") + for section in ( + "VULNERABILITY REPORT", + "WHAT IS IT?", + "WHY IS IT DANGEROUS?", + "VULNERABLE CODE:", + "FIXED CODE:", + "HOW TO FIX IT:", + ): + assert section in report + + +def test_format_report_handles_an_unknown_id(): + assert format_fix_report("CWE-999999").startswith("Not found:") + + +def test_main_exit_codes(capsys): + assert main(["119"]) == 0 + assert main(["CWE-999999"]) == 1 + assert main([]) == 2 + assert main(["--list"]) == 0 + + captured = capsys.readouterr() + assert "CWE-119" in captured.out + assert "No CWE id given" in captured.err + + +def test_database_path_is_independent_of_cwd(tmp_path, monkeypatch): + """Importing must work from any directory, not just ``src/models``.""" + module_name = "src.models.fix_suggester" + original = sys.modules[module_name] + + monkeypatch.chdir(tmp_path) + sys.modules.pop(module_name) + try: + reloaded = importlib.import_module(module_name) + assert reloaded.DATABASE_PATH.is_file() + assert reloaded.CWE_DATABASE + finally: + sys.modules[module_name] = original \ No newline at end of file diff --git a/tests/test_loader.py b/tests/test_loader.py new file mode 100644 index 0000000..5a5498a --- /dev/null +++ b/tests/test_loader.py @@ -0,0 +1,295 @@ +"""Tests for :mod:`src.models.loader`. + +The checkpoint validator exists because a ``BaselineMLP`` checkpoint used to be +loaded into a ``SecureScanModel`` with ``strict=False``. The two architectures +share no parameter names, so every weight was discarded and the model still +returned a confident probability. ``test_baseline_checkpoint_is_rejected`` is +the regression test for that. +""" + +from __future__ import annotations + +import torch +import pytest +from transformers import AutoConfig + +from src.models import loader +from src.models.baseline_mlp import BaselineMLP +from src.models.loader import ( + CheckpointError, + build_model, + diagnose_state_dict, + extract_state_dict, + load_checkpoint, + resolve_checkpoint, + save_checkpoint, +) +from src.models.securescan_model import SecureScanModel + + +# --------------------------------------------------------------------------- # +# resolve_checkpoint +# --------------------------------------------------------------------------- # + +def test_resolve_checkpoint_accepts_absolute_path(tmp_path): + path = tmp_path / "model.pt" + path.write_bytes(b"") + assert resolve_checkpoint(path) == path + + +def test_resolve_checkpoint_searches_the_legacy_directory(tmp_path, monkeypatch): + """A bare filename still resolves against the historical ``src/models``.""" + legacy = tmp_path / "src" / "models" + legacy.mkdir(parents=True) + target = legacy / "best_model.pt" + target.write_bytes(b"") + + monkeypatch.setattr(loader, "PROJECT_ROOT", tmp_path / "absent") + monkeypatch.setattr(loader, "DEFAULT_CHECKPOINT_DIR", tmp_path / "checkpoints") + monkeypatch.setattr(loader, "LEGACY_CHECKPOINT_DIR", legacy) + + assert resolve_checkpoint("best_model.pt") == target.resolve() + + +def test_resolve_checkpoint_ignores_the_working_directory(tmp_path, monkeypatch): + """Lookup is anchored to the project, not to wherever the script started.""" + project = tmp_path / "project" + (project / "checkpoints").mkdir(parents=True) + target = project / "checkpoints" / "best_model.pt" + target.write_bytes(b"") + + elsewhere = tmp_path / "elsewhere" + elsewhere.mkdir() + (elsewhere / "best_model.pt").write_bytes(b"a copy in the cwd") + + monkeypatch.setattr(loader, "PROJECT_ROOT", project) + monkeypatch.setattr(loader, "DEFAULT_CHECKPOINT_DIR", project / "checkpoints") + monkeypatch.setattr(loader, "LEGACY_CHECKPOINT_DIR", project / "src" / "models") + monkeypatch.chdir(elsewhere) + + assert resolve_checkpoint("best_model.pt") == target.resolve() + + +def test_resolve_checkpoint_reports_every_location_searched(tmp_path, monkeypatch): + monkeypatch.setattr(loader, "PROJECT_ROOT", tmp_path / "a") + monkeypatch.setattr(loader, "DEFAULT_CHECKPOINT_DIR", tmp_path / "b") + monkeypatch.setattr(loader, "LEGACY_CHECKPOINT_DIR", tmp_path / "c") + + with pytest.raises(FileNotFoundError) as excinfo: + resolve_checkpoint("nope.pt") + + message = str(excinfo.value) + assert "not found" in message + for location in ("a", "b", "c"): + assert str(tmp_path / location) in message + + +# --------------------------------------------------------------------------- # +# extract_state_dict +# --------------------------------------------------------------------------- # + +def test_extract_state_dict_passes_through_a_bare_state_dict(tiny_model): + state = tiny_model.state_dict() + assert extract_state_dict(state) is state + + +@pytest.mark.parametrize("key", ["model_state_dict", "state_dict"]) +def test_extract_state_dict_unwraps_provenance_wrappers(tiny_model, key): + """The wrapper written by ``src.utils.helpers.save_checkpoint``.""" + state = tiny_model.state_dict() + assert extract_state_dict({key: state, "epoch": 3}) is state + + +def test_extract_state_dict_accepts_a_whole_module(tiny_model): + extracted = extract_state_dict({"model": tiny_model}) + assert set(extracted) == set(tiny_model.state_dict()) + + +def test_extract_state_dict_rejects_unknown_formats(): + with pytest.raises(CheckpointError, match="Unrecognised checkpoint format"): + extract_state_dict({"epoch": 3, "val_f1": 0.9}) + + +# --------------------------------------------------------------------------- # +# diagnose_state_dict +# --------------------------------------------------------------------------- # + +def test_diagnose_reports_missing_and_unexpected_keys(tiny_model): + state = tiny_model.state_dict() + dropped = next(iter(state)) + partial = {k: v for k, v in state.items() if k != dropped} + partial["classifier.9.weight"] = torch.zeros(2) + + usable, missing, unexpected, shape_mismatch = diagnose_state_dict(tiny_model, partial) + + assert missing == [dropped] + assert unexpected == ["classifier.9.weight"] + assert shape_mismatch == [] + assert len(usable) == len(state) - 1 + + +def test_diagnose_rejects_shape_mismatches(tiny_model): + state = tiny_model.state_dict() + name = next(n for n, t in state.items() if t.dim() >= 2) + state[name] = torch.zeros((1,) + tuple(state[name].shape[1:])) + + usable, missing, unexpected, shape_mismatch = diagnose_state_dict(tiny_model, state) + + assert missing == [] and unexpected == [] + assert len(shape_mismatch) == 1 and name in shape_mismatch[0] + assert name not in usable + + +def test_diagnose_strips_data_parallel_prefix(tiny_model): + """A checkpoint saved through ``DataParallel`` maps onto the bare model.""" + own = tiny_model.state_dict() + wrapped = {f"module.{name}": tensor for name, tensor in own.items()} + + usable, missing, unexpected, shape_mismatch = diagnose_state_dict(tiny_model, wrapped) + + assert (missing, unexpected, shape_mismatch) == ([], [], []) + assert len(usable) == len(own) + + +# --------------------------------------------------------------------------- # +# save_checkpoint / load_checkpoint +# --------------------------------------------------------------------------- # + +def test_save_checkpoint_round_trips_every_parameter(make_tiny_model, tmp_path): + source = make_tiny_model() + original = {k: v.clone() for k, v in source.state_dict().items()} + destination = tmp_path / "nested" / "best_model.pt" + + written = save_checkpoint(source, destination, epoch=4, metrics={"val_f1": 0.9}) + + assert written == destination.resolve() + assert written.is_file() + + payload = torch.load(written, map_location="cpu", weights_only=True) + assert payload["epoch"] == 4 + assert payload["metrics"] == {"val_f1": 0.9} + + target = make_tiny_model() + load_checkpoint(target, written) + for name, tensor in target.state_dict().items(): + assert torch.equal(tensor, original[name]), name + + +def test_load_checkpoint_restores_from_a_helpers_wrapper(make_tiny_model, tmp_path): + """``src.utils.helpers.save_checkpoint`` writes ``model_state_dict``.""" + source = make_tiny_model() + torch.save( + { + "epoch": 1, + "model_state_dict": source.state_dict(), + "optimizer_state_dict": {}, + "val_f1": 0.5, + }, + tmp_path / "wrapped.pt", + ) + + target = make_tiny_model() + load_checkpoint(target, tmp_path / "wrapped.pt") + expected = source.state_dict() + for name, tensor in target.state_dict().items(): + assert torch.equal(tensor, expected[name]), name + + +def test_baseline_checkpoint_is_rejected(make_tiny_model, tmp_path): + """The bug this module exists to prevent. + + ``BaselineMLP`` and ``SecureScanModel`` share no parameter names, so the old + ``strict=False`` load discarded everything and left the classifier head + random, while the app kept reporting confident predictions. + """ + baseline_path = save_checkpoint(BaselineMLP(), tmp_path / "baseline.pt") + target = make_tiny_model() + before = {k: v.clone() for k, v in target.state_dict().items()} + + with pytest.raises(CheckpointError) as excinfo: + load_checkpoint(target, baseline_path) + + message = str(excinfo.value) + assert "does not fit" in message + assert "coverage" in message + # Refusing must mean refusing: nothing may be applied on the way out. + for name, tensor in target.state_dict().items(): + assert torch.equal(tensor, before[name]), name + + +def test_partial_load_must_be_requested_explicitly(tiny_model, tmp_path): + """``strict=False`` alone used to still demand 100% coverage.""" + state = dict(tiny_model.state_dict()) + del state[next(iter(state))] + torch.save({"model_state_dict": state}, tmp_path / "partial.pt") + + with pytest.warns(RuntimeWarning, match="Partial load"): + load_checkpoint(tiny_model, tmp_path / "partial.pt", strict=False) + + +def test_min_coverage_below_the_threshold_is_rejected(tiny_model, tmp_path): + state = dict(tiny_model.state_dict()) + del state[next(iter(state))] + torch.save({"model_state_dict": state}, tmp_path / "partial.pt") + + with pytest.raises(CheckpointError, match="covers only"): + load_checkpoint( + tiny_model, tmp_path / "partial.pt", strict=False, min_coverage=1.0 + ) + + +def test_min_coverage_outside_unit_range_is_rejected(tiny_model): + """Argument validation happens before any file access.""" + with pytest.raises(ValueError, match="min_coverage"): + load_checkpoint(tiny_model, "does-not-exist.pt", min_coverage=1.5) + + +# --------------------------------------------------------------------------- # +# build_model / from_config +# --------------------------------------------------------------------------- # + +def test_build_model_ignores_unknown_config_keys(): + model = build_model( + BaselineMLP, input_dim=16, hidden_dims=[4], dropout=0.0, not_a_kwarg=1 + ) + assert isinstance(model, BaselineMLP) + + +def test_build_model_applies_aliases(tiny_encoder_dir): + # The default encoder is 768 wide, so this only holds if the aliased key + # actually reached the constructor. + expected_width = AutoConfig.from_pretrained(str(tiny_encoder_dir)).hidden_size + + model = build_model( + SecureScanModel, + aliases={"codebert": "codebert_model"}, + codebert=str(tiny_encoder_dir), + freeze_layers=0, # the tiny encoder has fewer layers than the default 6 + from_scratch=True, + unknown_key="ignored", + ) + + assert isinstance(model, SecureScanModel) + assert model.codebert.config.hidden_size == expected_width + + +def test_from_config_accepts_the_config_file_spelling(tiny_encoder_dir): + """``phases/phase4-refinement/config.yaml`` spells the encoder ``codebert``.""" + expected_width = AutoConfig.from_pretrained(str(tiny_encoder_dir)).hidden_size + + model = SecureScanModel.from_config( + { + "name": "SecureScanModel", # unknown key, must be ignored + "codebert": str(tiny_encoder_dir), + "freeze_layers": 1, + "lstm_hidden": 8, + "lstm_layers": 1, + "dropout": 0.0, + "from_scratch": True, + } + ) + + assert model.codebert.config.hidden_size == expected_width + assert all( + not p.requires_grad for p in model.codebert.encoder.layer[0].parameters() + ) \ No newline at end of file diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 14a590d..6fabb02 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -1,17 +1,21 @@ -"""Smoke tests for preprocessing pipeline.""" +"""Smoke tests for the model architectures. + +The models are built from a local, tiny encoder config (see ``conftest.py``) so +these tests never download pretrained weights and stay fast enough for CI. +""" import torch import pytest -def test_model_output_shape(): + +def test_model_output_shape(tiny_model): """Test model outputs correct shape.""" - from src.models.securescan_model import SecureScanModel - model = SecureScanModel() - dummy_ids = torch.randint(0, 1000, (2, 512)) - dummy_mask = torch.ones(2, 512, dtype=torch.long) + dummy_ids = torch.randint(0, 64, (2, 32)) + dummy_mask = torch.ones(2, 32, dtype=torch.long) with torch.no_grad(): - output = model(dummy_ids, dummy_mask) + output = tiny_model(dummy_ids, dummy_mask) assert output.shape == (2, 2) + def test_baseline_mlp_shape(): """Test baseline MLP outputs correct shape.""" from src.models.baseline_mlp import BaselineMLP @@ -19,3 +23,19 @@ def test_baseline_mlp_shape(): dummy_input = torch.randn(4, 768) output = model(dummy_input) assert output.shape == (4, 2) + + +def test_freeze_layers_freezes_only_the_leading_layers(make_tiny_model): + """``freeze_layers`` freezes exactly the first N encoder layers.""" + model = make_tiny_model(freeze_layers=1) + + frozen = model.codebert.encoder.layer[0] + trainable = model.codebert.encoder.layer[1] + assert all(not p.requires_grad for p in frozen.parameters()) + assert all(p.requires_grad for p in trainable.parameters()) + + +def test_freeze_layers_above_layer_count_is_rejected(make_tiny_model): + """A freeze count larger than the encoder must fail loudly, not silently.""" + with pytest.raises(ValueError, match="freeze_layers"): + make_tiny_model(freeze_layers=3)