weightify 1.0.0__tar.gz

weightify-1.0.0/PKG-INFO

@@ -0,0 +1,12 @@
+Metadata-Version: 2.4
+Name: weightify
+Version: 1.0.0
+Summary: A lightweight security scanner for AI model weights.
+Project-URL: Homepage, https://github.com/Dur-E-Nayab-Khan/weightify
+Author-email: Dur E Nayab <durenayabkhan459@gmail.com>
+License: MIT
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Security
+Requires-Python: >=3.8

weightify-1.0.0/pyproject.toml

@@ -0,0 +1,22 @@
+# pyproject.toml
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "weightify"
+version = "1.0.0"
+description = "A lightweight security scanner for AI model weights."
+readme = "README.md"
+requires-python = ">=3.8"
+license = {text = "MIT"}
+authors = [{ name = "Dur E Nayab", email = "durenayabkhan459@gmail.com" }]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+    "Topic :: Security",
+]
+
+[project.urls]
+Homepage = "https://github.com/Dur-E-Nayab-Khan/weightify"

weightify-1.0.0/setup.py

@@ -0,0 +1,11 @@
+# setup.py
+from setuptools import setup, find_packages
+
+setup(
+    name="weightify",
+    version="0.1.0",
+    packages=find_packages(),
+    description="A security scanner for AI model weights",
+    author="Your Name",
+    python_requires=">=3.8",
+)

weightify-1.0.0/weightify/__init__.py

@@ -0,0 +1,23 @@
+# weightify/__init__.py
+import io
+import zipfile
+from .engine import WeightScanner
+
+def scan(byte_data: bytes):
+    """The main API for Weightify."""
+    scanner = WeightScanner()
+    results = []
+    
+    # 1. Check if it's a Zip (Standard for PyTorch .pt files)
+    stream = io.BytesIO(byte_data)
+    if zipfile.is_zipfile(stream):
+        with zipfile.ZipFile(stream, 'r') as z:
+            for file_info in z.infolist():
+                if file_info.filename.endswith(('.pkl', '.bin')) or "data.pkl" in file_info.filename:
+                    with z.open(file_info) as f:
+                        results.extend(scanner.scan_stream(f.read()))
+    else:
+        # 2. Otherwise scan as raw bytes
+        results.extend(scanner.scan_stream(byte_data))
+        
+    return results

weightify-1.0.0/weightify/archive_handler.py

@@ -0,0 +1,16 @@
+# weightify/archive_handler.py
+import zipfile
+import io
+
+def get_pickle_streams(byte_data):
+    """Detects if data is a zip (PyTorch) or raw pickle and yields streams."""
+    stream = io.BytesIO(byte_data)
+    
+    if zipfile.is_zipfile(stream):
+        with zipfile.ZipFile(stream, 'r') as z:
+            for file_name in z.namelist():
+                # PyTorch stores weights in 'data.pkl' or files ending in .pkl
+                if file_name.endswith('.pkl') or 'data.pkl' in file_name:
+                    yield z.read(file_name)
+    else:
+        yield byte_data

weightify-1.0.0/weightify/constants.py

@@ -0,0 +1,68 @@
+# weightify/constants.py
+
+"""
+Weightify Threat Database
+-------------------------
+This file contains the known dangerous Python globals that can be 
+used to achieve Remote Code Execution (RCE) via serialized model weights.
+"""
+
+# CRITICAL: These modules/functions trigger an immediate 'Dangerous' finding.
+# We use "*" to block an entire module if any function within it is a risk.
+UNSAFE_GLOBALS = {
+    # Core Python builtins that allow code execution or file manipulation
+    "__builtin__": {
+        "eval", "compile", "getattr", "apply", 
+        "exec", "open", "breakpoint", "input"
+    },
+    "builtins": {
+        "eval", "compile", "getattr", "apply", 
+        "exec", "open", "breakpoint", "input"
+    },
+
+    # System & Process Control (The most common RCE vectors)
+    "os": "*",            
+    "subprocess": "*",    
+    "posix": "*",         # os alias for Linux/Mac
+    "nt": "*",            # os alias for Windows
+    "pty": "*",           # Pseudo-terminal utilities
+    
+    # Network & Data Exfiltration
+    "socket": "*",
+    "requests.api": "*",
+    "urllib.request": "*",
+    "http.client": "*",
+    
+    # Advanced code reconstruction (used by hackers to hide logic)
+    "types": {"CodeType", "FunctionType", "ModuleType"},
+    "runpy": "*",
+    "code": "*",
+    
+    # File System manipulation
+    "shutil": "*",
+    "pathlib": "*",
+    "tempfile": "*",
+    
+    # Tools often used for obfuscation/persistence
+    "sys": "*",
+    "importlib": "*",
+    "pickle": "*",
+    "shelve": "*"
+}
+
+# INNOCUOUS: These are allowed math/data structures. 
+# Anything NOT in this list and NOT in the UNSAFE list is marked as 'Suspicious'.
+SAFE_GLOBALS = {
+    "numpy": {
+        "dtype", "ndarray", "core.multiarray._reconstruct", 
+        "_core.multiarray._reconstruct"
+    },
+    "torch": {
+        "FloatStorage", "LongStorage", "HalfStorage", "ByteStorage",
+        "DoubleStorage", "CharStorage", "ShortStorage", "IntStorage",
+        "BoolStorage", "BFloat16Storage", "ComplexFloatStorage",
+        "_utils._rebuild_tensor_v2"
+    },
+    "collections": {"OrderedDict"},
+    "builtins": {"set", "list", "dict", "tuple"}
+}

weightify-1.0.0/weightify/engine.py

@@ -0,0 +1,27 @@
+# weightify/engine.py
+import pickletools
+import io
+from .constants import UNSAFE_GLOBALS
+
+class WeightScanner:
+    def __init__(self):
+        self.blocklist = UNSAFE_GLOBALS
+
+    def scan_stream(self, data_stream):
+        """Analyzes a single byte stream for vulnerabilities."""
+        findings = []
+        try:
+            for opcode, arg, pos in pickletools.genops(data_stream):
+                # Standard Global calls
+                if opcode.name in ("GLOBAL", "INST"):
+                    module, name = arg.split(" ", 1)
+                    if self._is_unsafe(module, name):
+                        findings.append({"module": module, "op": name, "at": pos})
+        except Exception:
+            pass
+        return findings
+
+    def _is_unsafe(self, module, name):
+        if module in self.blocklist:
+            return self.blocklist[module] == "*" or name in self.blocklist[module]
+        return False

weightify-1.0.0/weightify.egg-info/PKG-INFO

@@ -0,0 +1,10 @@
+Metadata-Version: 2.4
+Name: weightify
+Version: 1.0.0
+Summary: A lightweight security scanner for AI model weights.
+Author: Your Name
+Author-email: Your Name <you@example.com>
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Dynamic: author
+Dynamic: requires-python

weightify-1.0.0/weightify.egg-info/SOURCES.txt

@@ -0,0 +1,12 @@
+README.md
+pyproject.toml
+setup.py
+weightify/__init__.py
+weightify/archive_handler.py
+weightify/constants.py
+weightify/engine.py
+weightify.egg-info/PKG-INFO
+weightify.egg-info/SOURCES.txt
+weightify.egg-info/dependency_links.txt
+weightify.egg-info/entry_points.txt
+weightify.egg-info/top_level.txt

weightify-1.0.0/weightify.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

weightify-1.0.0/weightify.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+weightify-scan = weightify.cli:main

weightify-1.0.0/weightify.egg-info/top_level.txt

@@ -0,0 +1 @@
+weightify