#!/usr/bin/env python3
"""
vulnerable_serverless.py - Simulated AWS Lambda-like runtime with event injection.
"""
import json, os, shutil, subprocess
from http.server import HTTPServer, BaseHTTPRequestHandler

FUNCTION_DIR = "/tmp/function"
os.makedirs(FUNCTION_DIR, exist_ok=True)
# initial function code
with open(os.path.join(FUNCTION_DIR, "handler.py"), "w") as f:
    f.write("""
def handler(event):
    return "Hello, " + event.get('name', 'world')
""")

class Handler(BaseHTTPRequestHandler):
    def do_POST(self):
        content_length = int(self.headers['Content-Length'])
        body = self.rfile.read(content_length)
        event = json.loads(body)
        # Vulnerable: use event['key'] to decide which file to load?
        # Simulate: the function source is overwritten by an "update" event from storage.
        if event.get('source') == 'storage':
            # Path traversal in object key
            object_key = event['object']['key']  # attacker controlled
            # Overwrite handler.py with the object content (simulated)
            dst = os.path.join(FUNCTION_DIR, "handler.py")
            # directory traversal to write outside? But we want to overwrite handler.py.
            # Attack: object.key = "../../../tmp/function/handler.py"
            # Normalize to ensure it's within FUNCTION_DIR? No validation!
            # The "get object" would fetch the file; here we just write injected code.
            injected_code = event.get('code', '# no code')
            # Resolve the full path – this is the vulnerability:
            full_path = os.path.normpath(os.path.join(FUNCTION_DIR, object_key))
            # Only check if it is under FUNCTION_DIR? Not present.
            with open(full_path, 'w') as f:
                f.write(injected_code)
            self.send_response(200)
            self.end_headers()
            self.wfile.write(b"Update applied")
        else:
            # Execute current handler (for demo)
            import handler
            result = handler.handler(event)
            self.send_response(200)
            self.end_headers()
            self.wfile.write(result.encode())

server = HTTPServer(('0.0.0.0', 8000), Handler)
server.serve_forever()