Update app.py

app.py

@@ -20,7 +20,12 @@ GITHUB_REPO_OWNER = "Wuhpondiscord"
 GITHUB_REPO_NAME = "ports"
 GITHUB_PORTS_DIR_PATH = "ports" # The directory within the repo where port files are located
 
-# --- Enhanced Port Parsing with Multi-Port Service Detection ---
+# Global variable to store parsed port definitions from all files
+# Format: {filename: {port_num: description}}
+ALL_GITHUB_PORT_DEFINITIONS: Dict[str, Dict[int, str]] = {}
+AVAILABLE_PORT_FILES: List[str] = [] # List of filenames from GitHub
+
+# --- GitHub File Listing and Caching ---
 
 def get_github_directory_contents(owner: str, repo: str, path: str) -> List[Dict]:
     """
@@ -37,120 +42,167 @@ def get_github_directory_contents(owner: str, repo: str, path: str) -> List[Dict
         print(f"⚠️ Error fetching GitHub directory listing for {path}: {e}")
         return []
 
-def load_custom_ports() -> Dict[int, str]:
+def parse_single_port_file_content(content: str) -> Dict[int, str]:
     """
-    Enhanced parser that:
-    1. Handles continuous format: '135 - MS SQL1433 - MSSQL1521 - Oracle'
-    2. Detects services on multiple ports
-    3. Groups related services
-    Loads port definitions from all .txt files in the specified GitHub directory.
+    Parses the content of a single port definition file.
+    Returns a dictionary of {port: description}.
     """
-    port_map = {80: "HTTP", 443: "HTTPS"}
-    service_ports = defaultdict(list)  # Track which ports belong to same service
-    
-    # Get the list of files from the GitHub directory
-    repo_contents = get_github_directory_contents(GITHUB_REPO_OWNER, GITHUB_REPO_NAME, GITHUB_PORTS_DIR_PATH)
-    
-    port_files_urls = []
-    for item in repo_contents:
-        if item.get("type") == "file" and item.get("name", "").endswith(('.txt', '.csv', '.conf', '.list')):
-            # The 'download_url' is the raw URL for the file content
-            port_files_urls.append((item['name'], item['download_url']))
-
-    if not port_files_urls:
-        print(f"⚠️ No port definition files found in '{GITHUB_PORTS_DIR_PATH}' on GitHub. Using default (80, 443).")
-        return port_map
+    single_file_port_map = {}
+    service_ports = defaultdict(list) # For multi-port detection within this file's context
 
-    print(f"✅ Found {len(port_files_urls)} port definition files to process.")
-
-    for filename, file_url in port_files_urls:
+    # Strategy 1: Continuous format parser
+    continuous_pattern = r'(\d+)\s*-\s*([A-Za-z][\w\s\-\.\(\)]*?)(?=\d+\s*-|$)'
+    matches = re.findall(continuous_pattern, content, re.MULTILINE)
+    
+    for port_str, desc in matches:
         try:
-            print(f"🌐 Fetching port definitions from: {filename} ({file_url})")
-            response = requests.get(file_url, timeout=10)
-            response.raise_for_status()
-            content = response.text
-            print(f"✅ Successfully fetched content for {filename}.")
-            
-            # --- Parsing Logic (remains largely the same) ---
-            
-            # Strategy 1: Continuous format parser
-            continuous_pattern = r'(\d+)\s*-\s*([A-Za-z][\w\s\-\.\(\)]*?)(?=\d+\s*-|$)'
-            matches = re.findall(continuous_pattern, content, re.MULTILINE)
-            
-            for port_str, desc in matches:
+            port_num = int(port_str)
+            if 1 <= port_num <= 65535:
+                desc = desc.strip().strip('-').strip()
+                desc = re.sub(r'\s+', ' ', desc)
+                if desc:
+                    base_service = desc.split('-')[0].strip()
+                    service_ports[base_service].append(port_num)
+                    single_file_port_map[port_num] = desc
+        except ValueError:
+            continue
+    
+    # Strategy 2: Line-by-line format
+    for line in content.split('\n'):
+        line = line.strip()
+        if not line or line.startswith('#'):
+            continue
+        
+        if re.match(r'^\d+\s*-\s*[A-Za-z]', line): # Skip if already matched by continuous pattern
+            continue
+        
+        patterns = [
+            r'^(.+?)\s*[:=]\s*(\d+)', # e.g., MySQL: 3306
+            r'port\s*[:=]?\s*(\d+)\s+(.+)', # e.g., port 22 SSH
+            r'(\d+)\s*/\s*\w+\s+(.+)', # e.g., 22/tcp SSH
+            r'^(\d+)\s+(.+)$', # e.g., 22 SSH
+        ]
+        
+        for pattern in patterns:
+            match = re.match(pattern, line, re.IGNORECASE)
+            if match:
+                groups = match.groups()
+                
+                if groups[0].isdigit():
+                    port_str, desc = groups[0], groups[1]
+                else:
+                    desc, port_str = groups[0], groups[1]
+                
                 try:
                     port_num = int(port_str)
                     if 1 <= port_num <= 65535:
-                        desc = desc.strip().strip('-').strip()
+                        desc = re.sub(r'[^\w\s\-\.]+', '', desc).strip()
                         desc = re.sub(r'\s+', ' ', desc)
                         if desc:
-                            # Extract base service name
                             base_service = desc.split('-')[0].strip()
                             service_ports[base_service].append(port_num)
-                            port_map[port_num] = desc
+                            single_file_port_map[port_num] = desc
+                        break
                 except ValueError:
                     continue
-            
-            # Strategy 2: Line-by-line format
-            for line in content.split('\n'):
-                line = line.strip()
-                if not line or line.startswith('#'):
-                    continue
-                
-                # Skip if already matched by continuous pattern
-                if re.match(r'^\d+\s*-\s*[A-Za-z]', line):
-                    continue
-                
-                patterns = [
-                    r'^(.+?)\s*[:=]\s*(\d+)', # e.g., MySQL: 3306
-                    r'port\s*[:=]?\s*(\d+)\s+(.+)', # e.g., port 22 SSH
-                    r'(\d+)\s*/\s*\w+\s+(.+)', # e.g., 22/tcp SSH
-                    r'^(\d+)\s+(.+)$', # e.g., 22 SSH
-                ]
-                
-                for pattern in patterns:
-                    match = re.match(pattern, line, re.IGNORECASE)
-                    if match:
-                        groups = match.groups()
-                        
-                        if groups[0].isdigit():
-                            port_str, desc = groups[0], groups[1]
-                        else:
-                            desc, port_str = groups[0], groups[1]
-                        
-                        try:
-                            port_num = int(port_str)
-                            if 1 <= port_num <= 65535:
-                                desc = re.sub(r'[^\w\s\-\.]+', '', desc).strip()
-                                desc = re.sub(r'\s+', ' ', desc)
-                                if desc:
-                                    base_service = desc.split('-')[0].strip()
-                                    service_ports[base_service].append(port_num)
-                                    port_map[port_num] = desc
-                                break
-                        except ValueError:
-                            continue
-                                
-        except requests.exceptions.RequestException as e:
-            print(f"⚠️ Error fetching {filename} from GitHub: {e}")
-        except Exception as e:
-            print(f"⚠️ Error parsing {filename}: {e}")
     
-    # Add multi-port info to descriptions
+    # Add multi-port info to descriptions (within this single file's context)
     for service, ports in service_ports.items():
         if len(ports) > 1:
             ports_sorted = sorted(ports)
             for port in ports_sorted:
-                if port in port_map:
-                    current_desc = port_map[port]
+                if port in single_file_port_map:
+                    current_desc = single_file_port_map[port]
                     if "also on" not in current_desc.lower():
                         other_ports = [str(p) for p in ports_sorted if p != port]
                         if other_ports:
-                            port_map[port] = f"{current_desc} (also on {','.join(other_ports[:3])}{'...' if len(other_ports) > 3 else ''})"
+                            single_file_port_map[port] = f"{current_desc} (also on {','.join(other_ports[:3])}{'...' if len(other_ports) > 3 else ''})"
+                            
+    return single_file_port_map
+
+def cache_all_github_port_files():
+    """
+    Fetches and parses all port definition files from GitHub, storing them in
+    the global ALL_GITHUB_PORT_DEFINITIONS and AVAILABLE_PORT_FILES.
+    """
+    global ALL_GITHUB_PORT_DEFINITIONS, AVAILABLE_PORT_FILES
+    ALL_GITHUB_PORT_DEFINITIONS.clear()
+    AVAILABLE_PORT_FILES.clear()
+
+    repo_contents = get_github_directory_contents(GITHUB_REPO_OWNER, GITHUB_REPO_NAME, GITHUB_PORTS_DIR_PATH)
     
-    print(f"✅ Loaded {len(port_map)} port definitions in total from GitHub.")
-    return port_map
+    port_files_to_fetch = []
+    for item in repo_contents:
+        if item.get("type") == "file" and item.get("name", "").endswith(('.txt', '.csv', '.conf', '.list')):
+            port_files_to_fetch.append((item['name'], item['download_url']))
+
+    if not port_files_to_fetch:
+        print(f"⚠️ No port definition files found in '{GITHUB_PORTS_DIR_PATH}' on GitHub.")
+        return
+
+    print(f"✅ Found {len(port_files_to_fetch)} port definition files to cache.")
+
+    for filename, file_url in port_files_to_fetch:
+        try:
+            print(f"🌐 Caching {filename} from {file_url}")
+            response = requests.get(file_url, timeout=10)
+            response.raise_for_status()
+            content = response.text
+            ALL_GITHUB_PORT_DEFINITIONS[filename] = parse_single_port_file_content(content)
+            AVAILABLE_PORT_FILES.append(filename)
+            print(f"✅ Cached {filename} with {len(ALL_GITHUB_PORT_DEFINITIONS[filename])} ports.")
+        except requests.exceptions.RequestException as e:
+            print(f"⚠️ Error fetching {filename} for caching: {e}")
+        except Exception as e:
+            print(f"⚠️ Error parsing {filename} for caching: {e}")
+
+    AVAILABLE_PORT_FILES.sort() # Keep the list sorted for UI
+    print(f"Finished caching. Total {len(AVAILABLE_PORT_FILES)} files available.")
+
+# --- Dynamic Port Map Construction ---
+
+def get_selected_port_map(selected_files: List[str]) -> Dict[int, str]:
+    """
+    Combines port definitions from selected files into a single port_map.
+    """
+    combined_port_map = {80: "HTTP", 443: "HTTPS"} # Always include defaults
+    all_service_ports = defaultdict(list) # To track multi-port across all selected files
 
+    if not selected_files:
+        print("No port files selected. Using default ports (80, 443).")
+        return combined_port_map
+
+    for filename in selected_files:
+        if filename in ALL_GITHUB_PORT_DEFINITIONS:
+            file_ports = ALL_GITHUB_PORT_DEFINITIONS[filename]
+            for port_num, desc in file_ports.items():
+                if port_num not in combined_port_map:
+                    combined_port_map[port_num] = desc
+                else:
+                    # If port exists, append new description if different, or choose more verbose
+                    existing_desc = combined_port_map[port_num]
+                    if desc != existing_desc and len(desc) > len(existing_desc):
+                        combined_port_map[port_num] = desc
+                
+                # Update service_ports for multi-port detection across combined map
+                base_service = desc.split('-')[0].strip()
+                if port_num not in all_service_ports[base_service]:
+                    all_service_ports[base_service].append(port_num)
+
+    # Re-apply multi-port info based on the combined map
+    for service, ports in all_service_ports.items():
+        if len(ports) > 1:
+            ports_sorted = sorted(ports)
+            for port in ports_sorted:
+                if port in combined_port_map:
+                    current_desc = combined_port_map[port]
+                    if "also on" not in current_desc.lower():
+                        other_ports = [str(p) for p in ports_sorted if p != port]
+                        if other_ports:
+                            combined_port_map[port] = f"{current_desc} (also on {','.join(other_ports[:3])}{'...' if len(other_ports) > 3 else ''})"
+
+    print(f"Generated port map from {len(selected_files)} selected files with {len(combined_port_map)} total port definitions.")
+    return combined_port_map
 
 # --- IP Extraction ---
 
@@ -299,15 +351,17 @@ def probe_target(ip: str, port: int, port_desc: str) -> Optional[Dict]:
                 pass
 
 
-def start_analysis(file_obj, max_threads: int = 50, progress=gr.Progress()):
+def start_analysis(file_obj, max_threads: int, selected_port_files: List[str], progress=gr.Progress()):
     """Main analysis with pure threading."""
     if file_obj is None:
         return pd.DataFrame([{"Error": "No file uploaded"}]), None
     
+    if not selected_port_files:
+        return pd.DataFrame([{"Error": "Please select at least one port list to scan."}]), None
+
     try:
-        progress(0, desc="📋 Parsing port definitions...")
-        # Call load_custom_ports without a folder_path argument
-        port_map = load_custom_ports() 
+        progress(0, desc="📋 Generating port map from selected lists...")
+        port_map = get_selected_port_map(selected_port_files)
         
         progress(0.1, desc="🔍 Extracting IP addresses...")
         ips = extract_ips(file_obj.name)
@@ -363,7 +417,7 @@ def start_analysis(file_obj, max_threads: int = 50, progress=gr.Progress()):
         return pd.DataFrame([{"Error": f"Analysis failed: {str(e)}"}]), None
 
 
-# --- UI ---
+# --- Gradio UI ---
 
 custom_css = """
 .gradio-container {
@@ -388,7 +442,7 @@ custom_css = """
 }
 """
 
-demo = gr.Blocks(title="C2 Deep Scanner Pro", css=custom_css) # Apply CSS here
+demo = gr.Blocks(title="C2 Deep Scanner Pro", css=custom_css)
 
 with demo:
     gr.HTML("""
@@ -398,22 +452,7 @@ with demo:
         </div>
     """)
     
-    with gr.Row():
-        gr.Markdown("""
-        ### 📌 Features
-        - **Enhanced Banner Detection**: HTTP, HTTPS, and raw socket banners
-        - **Multi-Port Services**: Shows related ports (e.g., MySQL on 3306, 3307)
-        - **Smart Protocol Detection**: Tries HTTP and HTTPS automatically
-        - **C2 Detection**: Cobalt Strike, Metasploit, Covenant, Empire & more
-        """)
-        
-        gr.Markdown("""
-        ### 📁 Input Files
-        - VirusTotal JSON exports
-        - SIEM/IDS log files
-        - Network scan results
-        - Plain text IP lists
-        """)
+    # Removed the large feature/input files markdown sections
     
     gr.Markdown("---")
     
@@ -440,25 +479,19 @@ with demo:
             )
         
         with gr.Column(scale=1):
-            gr.Markdown(f"""
-            ### 🔧 Port Format Support (from /{GITHUB_PORTS_DIR_PATH}/ on GitHub)
-            
-            **Continuous** (your format):
-            ```
-            135 - MS SQL1433 - MSSQL
-            ```
-            
-            **Standard formats**:
-            ```
-            80 - HTTP
-            MySQL: 3306
-            22/tcp SSH
-            ```
+            # Port List Selector
+            port_file_selector = gr.CheckboxGroup(
+                choices=AVAILABLE_PORT_FILES, # Will be populated on launch
+                value=AVAILABLE_PORT_FILES, # Select all by default
+                label=f"🔧 Select Port Lists (from /{GITHUB_PORTS_DIR_PATH}/ on GitHub)",
+                info="Choose which port definition files to use for scanning."
+            )
+            gr.Markdown("""
+            **Port Format Support:**
+            - **Continuous** (e.g., `135 - MS SQL1433 - MSSQL`)
+            - **Standard formats** (e.g., `80 - HTTP`, `MySQL: 3306`, `22/tcp SSH`)
             
-            Multi-port services shown as:
-            ```
-            MySQL (also on 3307,3308)
-            ```
+            Multi-port services shown as: `MySQL (also on 3307,3308)`
             """)
     
     gr.Markdown("---")
@@ -473,7 +506,7 @@ with demo:
     
     run_btn.click(
         fn=start_analysis,
-        inputs=[file_input, max_threads],
+        inputs=[file_input, max_threads, port_file_selector], # Added port_file_selector
         outputs=[output_table, download_btn]
     )
     
@@ -486,6 +519,13 @@ with demo:
 
 
 if __name__ == "__main__":
+    # Cache port files BEFORE launching the Gradio app
+    # This populates AVAILABLE_PORT_FILES for the CheckboxGroup
+    cache_all_github_port_files() 
+    
+    # Update the choices for the CheckboxGroup after caching
+    # This must be done here if the Gradio app is already defined
+    # or ensure AVAILABLE_PORT_FILES is populated before gr.CheckboxGroup is instantiated
     demo.queue(max_size=10)
     demo.launch(
         server_name="0.0.0.0",
@@ -493,5 +533,5 @@ if __name__ == "__main__":
         share=False,
         show_error=True,
         theme=gr.themes.Soft(),
-        ssr_mode=False  # CRITICAL: Disables SSR to prevent event loop errors
+        ssr_mode=False
     )