Update bot.py

bot.py

@@ -2,7 +2,6 @@ import os
 import io
 import zipfile
 import base64
-import random
 import tempfile
 from fastapi import FastAPI, Request
 import aiohttp
@@ -31,22 +30,17 @@ def xor_bytes(data: bytes, key: bytes) -> bytes:
 
 def encode_bytes(file_bytes: bytes, filename: str):
     """Encrypt + wrap into Python file."""
-    # Step 1: Zip the file
     zipped = make_zip_bytes(filename, file_bytes)
-    
-    # Step 2: Encrypt with XOR
     key = rand_bytes(16)
     enc = xor_bytes(zipped, key)
     enc_b64 = base64.b64encode(enc).decode()
     key_b64 = base64.b64encode(key).decode()
     
-    # Step 3: Build output Python file
     wrapper = f"""# Auto Encrypted by JerryCoder Encryptor
-import base64,zlib
+import base64,zipfile,io,tempfile,os,runpy
 D=base64.b64decode('{enc_b64}')
 K=base64.b64decode('{key_b64}')
 O=bytes([b^K[i%len(K)] for i,b in enumerate(D)])
-import io,zipfile,tempfile,os,runpy
 t=tempfile.mkdtemp()
 zipfile.ZipFile(io.BytesIO(O)).extractall(t)
 p=[f for f in os.listdir(t) if f.endswith('.py')][0]
@@ -66,6 +60,10 @@ async def encode_endpoint(request: Request):
     file_url = data.get("url")
     filename = data.get("filename", "file.py")
 
+    # ✅ Only allow Python files
+    if not filename.endswith(".py"):
+        return {"ok": False, "error": "Only Python (.py) files are allowed"}
+
     if not file_url:
         return {"ok": False, "error": "No file URL provided"}
 
@@ -81,9 +79,7 @@ async def encode_endpoint(request: Request):
     with open(out_path, "wb") as f:
         f.write(encoded_bytes)
 
-    # ⚠️ NOTE: Hugging Face doesn’t auto-host tmp files publicly.
-    # You’ll need to upload to a hosting service (like tmpfiles.org, catbox.moe, etc.)
-    # For now, just return ok=True and filename.
+    # ⚠️ NOTE: Replace this URL with an actual file hosting if you need public download
     return {
         "ok": True,
         "result_url": f"https://jerrycoder-kuttuxd.hf.space/tmp/{out_name}",