Update app.py

app.py

@@ -1,121 +1,158 @@
-import gradio as gr
-import tempfile
-import os
-import subprocess
-from androguard.misc import AnalyzeAPK
-from transformers import pipeline, Conversation
-from sentence_transformers import SentenceTransformer, util
-import json
-
-# Global storage
-apk_context = {"smali": {}, "java": {}, "info": ""}
-
-# Step 1: Decompile APK
-def decompile_apk(apk_file):
-    if apk_file is None:
-        return "No file uploaded. Please upload an APK file."
-
-    temp_apk_path = None
-    output_dir = tempfile.mkdtemp()
-    try:
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".apk") as temp_apk:
-            temp_apk.write(apk_file.read())
-            temp_apk_path = temp_apk.name
-
-        # Decompile APK to Smali
-        smali_output = os.path.join(output_dir, "smali")
-        subprocess.run(["baksmali", "d", temp_apk_path, "-o", smali_output], check=True)
-
-        # Decompile APK to Java
-        java_output = os.path.join(output_dir, "java")
-        subprocess.run(["jadx", "-d", java_output, temp_apk_path], check=True)
-
-        # Load Smali and Java content
-        smali_files = {}
-        for root, _, files in os.walk(smali_output):
-            for file in files:
-                if file.endswith(".smali"):
-                    with open(os.path.join(root, file), "r") as f:
-                        smali_files[file] = f.read()
-
-        java_files = {}
-        for root, _, files in os.walk(java_output):
-            for file in files:
-                if file.endswith(".java"):
-                    with open(os.path.join(root, file), "r") as f:
-                        java_files[file] = f.read()
-
-        # Save to global context
-        apk_context["smali"] = smali_files
-        apk_context["java"] = java_files
-
-        return f"Decompilation successful. Extracted {len(smali_files)} Smali files and {len(java_files)} Java files."
-
-    except Exception as e:
-        return f"Error during decompilation: {str(e)}"
-    finally:
-        if temp_apk_path and os.path.exists(temp_apk_path):
-            os.unlink(temp_apk_path)
-
-# Step 2: Build Searchable Database
-def build_search_index():
-    smali_texts = [f"{k}\n{v}" for k, v in apk_context["smali"].items()]
-    java_texts = [f"{k}\n{v}" for k, v in apk_context["java"].items()]
-
-    # Use SentenceTransformer for embeddings
-    model = SentenceTransformer('all-MiniLM-L6-v2')
-
-    smali_embeddings = model.encode(smali_texts, convert_to_tensor=True)
-    java_embeddings = model.encode(java_texts, convert_to_tensor=True)
-
-    return model, smali_texts, smali_embeddings, java_texts, java_embeddings
-
-# Step 3: Query Chatbot
-def query_apk_chat(user_message):
-    if not apk_context["smali"] and not apk_context["java"]:
-        return "No decompiled APK available. Please upload and decompile an APK first."
-
-    try:
-        # Build the search index
-        model, smali_texts, smali_embeddings, java_texts, java_embeddings = build_search_index()
-
-        # Encode user message
-        query_embedding = model.encode(user_message, convert_to_tensor=True)
-
-        # Search Smali and Java
-        smali_scores = util.pytorch_cos_sim(query_embedding, smali_embeddings).squeeze(0)
-        java_scores = util.pytorch_cos_sim(query_embedding, java_embeddings).squeeze(0)
-
-        # Get top matches
-        smali_result = smali_texts[smali_scores.argmax().item()]
-        java_result = java_texts[java_scores.argmax().item()]
-
-        # Combine results into a response
-        response = f"**Relevant Smali Code:**\n\n{smali_result[:1000]}\n\n"
-        response += f"**Relevant Java Code:**\n\n{java_result[:1000]}"
-
-        return response
-
-    except Exception as e:
-        return f"Error during search: {str(e)}"
-
-# Gradio Interface
-apk_upload_interface = gr.Interface(
-    fn=decompile_apk,
-    inputs=gr.File(label="Upload APK File", file_types=['.apk']),
-    outputs="text",
-    title="APK Analyzer with Full Source Access",
-    description="Upload an APK to decompile it into Smali and Java source code."
-)
-
-chat_interface = gr.Interface(
-    fn=query_apk_chat,
-    inputs="text",
-    outputs="text",
-    title="APK Chatbot",
-    description="Ask questions about the APK's source code in Smali or Java."
-)
-
-# Combine interfaces into tabs
-iface = gr.TabbedInterface([apk_upload_interface, chat_interface], ["Upload & Decompile", "Chat with Source"])
-iface.launch()
+import os
+import tempfile
+import subprocess
+from androguard.misc import AnalyzeAPK
+from transformers import pipeline, Conversation
+from sentence_transformers import SentenceTransformer, util
+import gradio as gr
+
+# Global storage for APK data
+apk_context = {"smali": {}, "java": {}, "info": ""}
+
+# Function to install Baksmali and JADX
+def install_tools():
+    baksmali_path = "/usr/local/bin/baksmali.jar"
+    jadx_path = "/usr/local/bin/jadx"
+    
+    # Install Baksmali if not present
+    if not os.path.exists(baksmali_path):
+        print("Installing Baksmali...")
+        subprocess.run(
+            [
+                "curl",
+                "-L",
+                "-o",
+                baksmali_path,
+                "https://bitbucket.org/JesusFreke/smali/downloads/baksmali-2.5.2.jar"
+            ],
+            check=True
+        )
+    
+    # Install JADX if not present
+    if not os.path.exists(jadx_path):
+        print("Installing JADX...")
+        subprocess.run(
+            [
+                "curl",
+                "-L",
+                "-o",
+                "/usr/local/bin/jadx.zip",
+                "https://github.com/skylot/jadx/releases/download/v1.4.7/jadx-1.4.7.zip"
+            ],
+            check=True
+        )
+        subprocess.run(["unzip", "-o", "/usr/local/bin/jadx.zip", "-d", "/usr/local/bin"], check=True)
+        subprocess.run(["chmod", "+x", "/usr/local/bin/jadx/bin/jadx"], check=True)
+
+# Install tools at the start of the app
+install_tools()
+
+# Step 1: Decompile APK
+def decompile_apk(apk_file):
+    if apk_file is None:
+        return "No file uploaded. Please upload an APK file."
+    
+    temp_apk_path = None
+    output_dir = tempfile.mkdtemp()
+    try:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".apk") as temp_apk:
+            temp_apk.write(apk_file.read())
+            temp_apk_path = temp_apk.name
+
+        # Decompile APK to Smali
+        smali_output = os.path.join(output_dir, "smali")
+        subprocess.run(["java", "-jar", "/usr/local/bin/baksmali.jar", "d", temp_apk_path, "-o", smali_output], check=True)
+
+        # Decompile APK to Java
+        java_output = os.path.join(output_dir, "java")
+        subprocess.run(["/usr/local/bin/jadx/bin/jadx", "-d", java_output, temp_apk_path], check=True)
+
+        # Load Smali and Java content
+        smali_files = {}
+        for root, _, files in os.walk(smali_output):
+            for file in files:
+                if file.endswith(".smali"):
+                    with open(os.path.join(root, file), "r") as f:
+                        smali_files[file] = f.read()
+
+        java_files = {}
+        for root, _, files in os.walk(java_output):
+            for file in files:
+                if file.endswith(".java"):
+                    with open(os.path.join(root, file), "r") as f:
+                        java_files[file] = f.read()
+
+        # Save to global context
+        apk_context["smali"] = smali_files
+        apk_context["java"] = java_files
+
+        return f"Decompilation successful. Extracted {len(smali_files)} Smali files and {len(java_files)} Java files."
+
+    except Exception as e:
+        return f"Error during decompilation: {str(e)}"
+    finally:
+        if temp_apk_path and os.path.exists(temp_apk_path):
+            os.unlink(temp_apk_path)
+
+# Step 2: Build Searchable Database
+def build_search_index():
+    smali_texts = [f"{k}\n{v}" for k, v in apk_context["smali"].items()]
+    java_texts = [f"{k}\n{v}" for k, v in apk_context["java"].items()]
+
+    # Use SentenceTransformer for embeddings
+    model = SentenceTransformer('all-MiniLM-L6-v2')
+
+    smali_embeddings = model.encode(smali_texts, convert_to_tensor=True)
+    java_embeddings = model.encode(java_texts, convert_to_tensor=True)
+
+    return model, smali_texts, smali_embeddings, java_texts, java_embeddings
+
+# Step 3: Query Chatbot
+def query_apk_chat(user_message):
+    if not apk_context["smali"] and not apk_context["java"]:
+        return "No decompiled APK available. Please upload and decompile an APK first."
+
+    try:
+        # Build the search index
+        model, smali_texts, smali_embeddings, java_texts, java_embeddings = build_search_index()
+
+        # Encode user message
+        query_embedding = model.encode(user_message, convert_to_tensor=True)
+
+        # Search Smali and Java
+        smali_scores = util.pytorch_cos_sim(query_embedding, smali_embeddings).squeeze(0)
+        java_scores = util.pytorch_cos_sim(query_embedding, java_embeddings).squeeze(0)
+
+        # Get top matches
+        smali_result = smali_texts[smali_scores.argmax().item()]
+        java_result = java_texts[java_scores.argmax().item()]
+
+        # Combine results into a response
+        response = f"**Relevant Smali Code:**\n\n{smali_result[:1000]}\n\n"
+        response += f"**Relevant Java Code:**\n\n{java_result[:1000]}"
+
+        return response
+
+    except Exception as e:
+        return f"Error during search: {str(e)}"
+
+# Gradio Interface
+apk_upload_interface = gr.Interface(
+    fn=decompile_apk,
+    inputs=gr.File(label="Upload APK File", file_types=['.apk']),
+    outputs="text",
+    title="APK Analyzer with Full Source Access",
+    description="Upload an APK to decompile it into Smali and Java source code."
+)
+
+chat_interface = gr.Interface(
+    fn=query_apk_chat,
+    inputs="text",
+    outputs="text",
+    title="APK Chatbot",
+    description="Ask questions about the APK's source code in Smali or Java."
+)
+
+# Combine interfaces into tabs
+iface = gr.TabbedInterface([apk_upload_interface, chat_interface], ["Upload & Decompile", "Chat with Source"])
+iface.launch()