T09 · Insecure Skill Coding Practices
Error
- Location
- scripts/predict.py:84
- Finding
- Arbitrary Code Execution Through Unsafe PyTorch Checkpoint Deserialization<![CDATA[ ## Vulnerability Details **File Location**: `scripts/predict.py:84` and `scripts/predict.py:141` **Vulnerability Type**: Unsafe deserialization of attacker-controlled PyTorch checkpoints **Risk Level**: High ### Vulnerable Code ```python def predict_gene(base_dir, genotype_df, location_code, device): model_path = os.path.join(base_dir, "data", "models_gene", f"{location_code}.pt") if not os.path.exists(model_path): return None ckpt = torch.load(model_path, map_location=device, weights_only=False) ``` ```python def predict_env(base_dir, genotype_df, env_features_normalized, device, traits=None): # ... for trait_code in traits: model_path = os.path.join(base_dir, "data", "models_env", f"{trait_code}.pt") if not os.path.exists(model_path): continue ckpt = torch.load(model_path, map_location=device, weights_only=False) ``` The checkpoint root is selected through a command-line argument: ```python parser.add_argument("--base_dir", default=None, help="Path to rice_prediction directory (auto-detected if omitted)") ``` ### Technical Analysis `torch.load()` uses Python pickle-compatible deserialization for checkpoint objects. Explicitly setting `weights_only=False` permits arbitrary Python objects in a checkpoint to be reconstructed. A malicious pickle object can define a reduction routine that invokes an operating-system command during deserialization. The caller can control the checkpoint source directory through `--base_dir`. The application performs only an existence check before loading the file. It does not enforce use of the canonical Skill directory, validate a cryptographic digest, verify a signature, or otherwise establish that the checkpoint is trusted. Code execution occurs while `torch.load()` is processing the checkpoint, before accesses such as `ckpt["model_state"]`. A malicious checkpoint therefore does not need to contain a valid model state or scale ...[truncated 1611 chars]
- Remediation
- <![CDATA[ ## Remediation Suggestions 1. Load tensor-only checkpoints with the restrictive mode: ```python ckpt = torch.load(model_path, map_location=device, weights_only=True) ``` 2. Do not serialize scikit-learn scaler objects inside pickle-based checkpoints. Store only primitive numeric scaler parameters, such as means and scales, in a non-executable format such as JSON, NPZ, or safetensors. 3. Prefer safetensors for model weights because it does not execute Python object reconstruction. 4. Publish SHA-256 digests or signed manifests for every model and verify them before loading. 5. Resolve the model path canonically and require it to remain inside the installed Skill's trusted model directory. 6. Remove `--base_dir` if arbitrary model roots are not required. If it is required, clearly treat it as a trusted-administrator option rather than ordinary user input. 7. Reject symbolic links and unexpected file types when validating model files. 8. Update `check_env.py` so its integrity check validates model hashes rather than checking only whether files exist. ]]>