Vulnerable code sample

import torch
import os
import argparse

def load_synthesis_model(model_dir):
    """
    Simulates loading a model during server startup.
    """
    # In the original code, a model file is loaded from a user-specified directory.
    model_path = os.path.join(model_dir, 'model.pth')
    
    print(f"[*] Attempting to load model from: {model_path}")

    # The vulnerable function call.
    # It uses torch.load() on a file from a user-controlled path
    # without specifying `weights_only=True`. This defaults to using
    # Python's pickle module, allowing for arbitrary code execution if
    # the .pth file is maliciously crafted.
    model = torch.load(model_path)

    print("[+] Model loaded successfully.")
    return model

def main():
    """
    Simulates the main entry point of the gRPC server application.
    """
    parser = argparse.ArgumentParser(
        description="Vulnerable server startup example."
    )
    parser.add_argument(
        "--model_dir",
        required=True,
        help="Path to the directory containing model files."
    )
    args = parser.parse_args()

    print("[*] Initializing server...")
    try:
        # The vulnerable model loading is called during the server's initialization phase.
        load_synthesis_model(args.model_dir)
        print("[+] Server started.")
    except Exception as e:
        print(f"[!] Server initialization failed: {e}")

if __name__ == "__main__":
    main()