Scan models for security risks

import sys

import hiddenlayer
from hiddenlayer import HiddenLayer

MODEL_PATH = "models/model.pkl"  # Path to the file you want to scan
BLOCKING_SEVERITIES = {"critical", "high"}


def main() -> int:
    # The client reads HIDDENLAYER_CLIENT_ID and HIDDENLAYER_CLIENT_SECRET
    # (or HIDDENLAYER_TOKEN) from the environment.
    with HiddenLayer() as client:
        try:
            report = client.model_scanner.scan_file(
                model_name="my-model",
                model_version="1",
                model_path=MODEL_PATH,
            )
        except FileNotFoundError:
            print(f"Model file not found: {MODEL_PATH}", file=sys.stderr)
            return 2
        except hiddenlayer.AuthenticationError:
            print("Authentication failed. Check your client ID, client secret, and environment.", file=sys.stderr)
            return 2
        except hiddenlayer.APIStatusError as e:
            print(f"HiddenLayer API error {e.status_code}: {e.message}", file=sys.stderr)
            return 2
        except hiddenlayer.APIConnectionError as e:
            print(f"Could not reach HiddenLayer: {e}", file=sys.stderr)
            return 2

    print(f"Scan ID:          {report.scan_id}")
    print(f"Status:           {report.status}")
    print(f"Highest severity: {report.summary.highest_severity}")
    print(f"Detections:       {report.summary.detection_count}")

    for file_result in report.file_results or []:
        print(f"\nFile: {file_result.file_location} ({file_result.status})")
        for detection in file_result.detections:
            print(f"  [{detection.severity}] {detection.category}: {detection.description}")

    if report.status != "done":
        print(f"Scan did not complete (status: {report.status}).", file=sys.stderr)
        return 2

    if report.summary.highest_severity in BLOCKING_SEVERITIES:
        print("Blocking: model has critical or high severity detections.", file=sys.stderr)
        return 1

    return 0


if __name__ == "__main__":
    sys.exit(main())

import sys

import hiddenlayer
from hiddenlayer import HiddenLayer

client = HiddenLayer()

try:
    report = client.model_scanner.scan_folder(
        model_name="my-folder-model",
        model_version="1",
        path="models/my-model",
        allow_file_patterns=["*.bin", "*.pt", "*.safetensors", "*.json"],
        ignore_file_patterns=["*.png"],
    )
except hiddenlayer.APIStatusError as e:
    sys.exit(f"HiddenLayer API error {e.status_code}: {e.message}")
except hiddenlayer.APIConnectionError as e:
    sys.exit(f"Could not reach HiddenLayer: {e}")
finally:
    client.close()

print(f"{report.status}: {report.summary.file_count} files, "
      f"{report.summary.detection_count} detections, "
      f"highest severity {report.summary.highest_severity}")

import os
import sys

import hiddenlayer
from hiddenlayer import HiddenLayer

client = HiddenLayer()

try:
    report = client.model_scanner.scan_huggingface_model(
        repo_id="openai-community/gpt2",
        # Optional: branch, tag, or commit hash. Also used as the model version.
        # revision="main",
        allow_file_patterns=["*.safetensors", "*.bin", "*.json"],
        # Use HF_TOKEN if set; otherwise read the token saved by the Hugging Face CLI.
        hf_token=os.environ.get("HF_TOKEN") or True,
    )
except ImportError as e:
    sys.exit(f"{e} Run: pip install huggingface_hub")
except hiddenlayer.APIStatusError as e:
    sys.exit(f"HiddenLayer API error {e.status_code}: {e.message}")
except hiddenlayer.APIConnectionError as e:
    sys.exit(f"Could not reach HiddenLayer: {e}")
finally:
    client.close()

print(f"{report.inventory.model_name}: {report.status}, "
      f"highest severity {report.summary.highest_severity}")

import sys

import hiddenlayer
from hiddenlayer import HiddenLayer

client = HiddenLayer()

try:
    report = client.model_scanner.scan_s3_model(
        model_name="my-s3-model",
        bucket="my-model-bucket",
        key="models/v1/model.pkl",
    )
except ImportError as e:
    sys.exit(f"{e} Run: pip install boto3")
except RuntimeError as e:
    # Raised when the download from S3 fails, for example due to missing AWS permissions.
    sys.exit(str(e))
except hiddenlayer.APIStatusError as e:
    sys.exit(f"HiddenLayer API error {e.status_code}: {e.message}")
except hiddenlayer.APIConnectionError as e:
    sys.exit(f"Could not reach HiddenLayer: {e}")
finally:
    client.close()

print(f"{report.status}: highest severity {report.summary.highest_severity}")

import sys

import hiddenlayer
from hiddenlayer import HiddenLayer
from hiddenlayer.lib import CommunityScanSource

client = HiddenLayer()

try:
    report = client.community_scanner.community_scan(
        model_name="bert-tiny-torch-vuln",
        model_path="drhyrum/bert-tiny-torch-vuln",  # Hugging Face repository ID
        model_source=CommunityScanSource.HUGGING_FACE,
        model_version="main",  # Default is "main"
    )
except ValueError as e:
    # Raised if the API does not return a scan ID.
    sys.exit(str(e))
except hiddenlayer.APIStatusError as e:
    sys.exit(f"HiddenLayer API error {e.status_code}: {e.message}")
except hiddenlayer.APIConnectionError as e:
    sys.exit(f"Could not reach HiddenLayer: {e}")
finally:
    client.close()

print(f"{report.scan_id}: {report.status}, "
      f"{report.summary.detection_count} detections, "
      f"highest severity {report.summary.highest_severity}")

from hiddenlayer import HiddenLayer

with HiddenLayer() as client:
    started = client.model_scanner.scan_file(
        model_name="my-model",
        model_path="models/model.pkl",
        wait_for_results=False,
    )
    print(f"Started scan {started.scan_id} (status: {started.status})")

    # Later, possibly in another process:
    report = client.model_scanner.scan_result(scan_id=started.scan_id)
    print(f"Status now: {report.status}")

import sys
import time

import hiddenlayer
from hiddenlayer import HiddenLayer

TIMEOUT_SECONDS = 15 * 60
POLL_INTERVAL_SECONDS = 10
TERMINAL_STATUSES = {"done", "failed", "canceled"}

with HiddenLayer() as client:
    try:
        started = client.model_scanner.scan_file(
            model_name="my-model",
            model_path="models/model.pkl",
            wait_for_results=False,
        )
        scan_id = started.scan_id
        deadline = time.monotonic() + TIMEOUT_SECONDS

        while True:
            summary = client.scans.results.retrieve_summary(scan_id)
            if summary.status in TERMINAL_STATUSES:
                break
            if time.monotonic() > deadline:
                sys.exit(f"Scan {scan_id} did not finish within {TIMEOUT_SECONDS} seconds "
                         f"(last status: {summary.status}).")
            time.sleep(POLL_INTERVAL_SECONDS)

        # Fetch the full report, including per-file results.
        report = client.model_scanner.scan_result(scan_id=scan_id)
    except hiddenlayer.APIStatusError as e:
        sys.exit(f"HiddenLayer API error {e.status_code}: {e.message}")
    except hiddenlayer.APIConnectionError as e:
        sys.exit(f"Could not reach HiddenLayer: {e}")

print(f"{report.status}: highest severity {report.summary.highest_severity}")