Improved mongodb, mysql and postgres scan

rohitcoder · rohitcoder · commit 4d573208bc33 · 2023-12-05T22:00:36.000+05:30
diff --git a/connection.yml.sample b/connection.yml.sample
@@ -36,14 +36,23 @@ sources:
       port: YOUR_MYSQL_PORT
       user: YOUR_MYSQL_USERNAME
       password: YOUR_MYSQL_PASSWORD
-      database: YOUR_MYSQL_DATABASE_NAME
+      limit_start: 0   # Specify the starting limit for the range
+      limit_end: 500   # Specify the ending limit for the range
+      tables:
+        - table1
+        - table2
   postgresql:
     postgresql_example:
       host: YOUR_POSTGRESQL_HOST
       port: YOUR_POSTGRESQL_PORT
       user: YOUR_POSTGRESQL_USERNAME
       password: YOUR_POSTGRESQL_PASSWORD
       database: YOUR_POSTGRESQL_DATABASE_NAME
+      limit_start: 0   # Specify the starting limit for the range
+      limit_end: 500   # Specify the ending limit for the range
+      tables:
+        - table1
+        - table2
   mongodb:
     mongodb_example:
       uri: YOUR_MONGODB_URI
@@ -52,6 +61,12 @@ sources:
       username: YOUR_MONGODB_USERNAME
       password: YOUR_MONGODB_PASSWORD
       database: YOUR_MONGODB_DATABASE_NAME
+      uri: YOUR_MONGODB_URI  # Use either URI or individual connection parameters
+      limit_start: 0   # Specify the starting limit for the range
+      limit_end: 500   # Specify the ending limit for the range
+      collections:
+        - collection1
+        - collection2
   fs:
     fs_example:
       path: /path/to/your/filesystem/directory
diff --git a/hawk_scanner/commands/couchdb.py b/hawk_scanner/commands/couchdb.py
@@ -0,0 +1,77 @@
+import couchdb
+from hawk_scanner.internals import system
+from rich.console import Console
+from rich.table import Table
+
+console = Console()
+
+def connect_couchdb(host, port, username, password, database):
+    try:
+        server = couchdb.Server(f"http://{username}:{password}@{host}:{port}/")
+        if database not in server:
+            system.print_error(f"Database {database} not found on CouchDB server.")
+            return None
+        db = server[database]
+        system.print_info(f"Connected to CouchDB database")
+        return db
+    except Exception as e:
+        system.print_error(f"Failed to connect to CouchDB database with error: {e}")
+        return None
+
+def check_data_patterns(db, patterns, profile_name, database_name):
+    results = []
+    for doc_id in db:
+        document = db[doc_id]
+        for field_name, field_value in document.items():
+            if field_value:
+                value_str = str(field_value)
+                matches = system.match_strings(value_str)
+                if matches:
+                    for match in matches:
+                        results.append({
+                            'host': f"{db.resource.credentials[1]}:{db.resource.credentials[2]}",
+                            'database': database_name,
+                            'document_id': doc_id,
+                            'field': field_name,
+                            'pattern_name': match['pattern_name'],
+                            'matches': match['matches'],
+                            'sample_text': match['sample_text'],
+                            'profile': profile_name,
+                            'data_source': 'couchdb'
+                        })
+
+    return results
+
+def execute(args):
+    results = []
+    system.print_info(f"Running Checks for CouchDB Sources")
+    connections = system.get_connection()
+
+    if 'sources' in connections:
+        sources_config = connections['sources']
+        couchdb_config = sources_config.get('couchdb')
+
+        if couchdb_config:
+            patterns = system.get_fingerprint_file()
+
+            for key, config in couchdb_config.items():
+                host = config.get('host')
+                port = config.get('port', 5984)  # default CouchDB port
+                username = config.get('username')
+                password = config.get('password')
+                database = config.get('database')
+
+                if host and username and password and database:
+                    system.print_info(f"Checking CouchDB Profile {key} with host and authentication")
+                else:
+                    system.print_error(f"Incomplete CouchDB configuration for key: {key}")
+                    continue
+
+                db = connect_couchdb(host, port, username, password, database)
+                if db:
+                    results += check_data_patterns(db, patterns, key, database)
+        else:
+            system.print_error("No CouchDB connection details found in connection.yml")
+    else:
+        system.print_error("No 'sources' section found in connection.yml")
+    return results
diff --git a/hawk_scanner/commands/mongodb.py b/hawk_scanner/commands/mongodb.py
@@ -1,8 +1,6 @@
 import pymongo
 from hawk_scanner.internals import system
-import re
 from rich.console import Console
-from rich.table import Table
 
 console = Console()
 
@@ -25,11 +23,11 @@ def connect_mongodb(host, port, username, password, database, uri=None):
         return None
 
 
-def check_data_patterns(db, patterns, profile_name, database_name):
+def check_data_patterns(db, patterns, profile_name, database_name, limit_start=0, limit_end=500):
     results = []
     for collection_name in db.list_collection_names():
         collection = db[collection_name]
-        for document in collection.find():
+        for document in collection.find().limit(limit_end).skip(limit_start):
             for field_name, field_value in document.items():
                 if field_value:
                     value_str = str(field_value)
@@ -69,6 +67,8 @@ def execute(args):
                 password = config.get('password')
                 database = config.get('database')
                 uri = config.get('uri')  # Added support for URI
+                limit_start = config.get('limit_start', 0)
+                limit_end = config.get('limit_end', 500)
 
                 if uri:
                     system.print_info(f"Checking MongoDB Profile {key} using URI")
@@ -80,7 +80,7 @@ def execute(args):
 
                 db = connect_mongodb(host, port, username, password, database, uri)
                 if db:
-                    results += check_data_patterns(db, patterns, key, database)
+                    results += check_data_patterns(db, patterns, key, database, limit_start=limit_start, limit_end=limit_end)
         else:
             system.print_error("No MongoDB connection details found in connection.yml")
     else:
diff --git a/hawk_scanner/commands/mysql.py b/hawk_scanner/commands/mysql.py
@@ -1,36 +1,37 @@
-import psycopg2
+import pymysql
 from hawk_scanner.internals import system
-import re
 from rich.console import Console
-from rich.table import Table
 
 console = Console()
 
-def connect_postgresql(host, port, user, password, database):
+def connect_mysql(host, port, user, password, database):
     try:
-        conn = psycopg2.connect(
+        conn = pymysql.connect(
             host=host,
             port=port,
             user=user,
             password=password,
             database=database
         )
         if conn:
-            system.print_info(f"Connected to PostgreSQL database at {host}")
+            system.print_info(f"Connected to MySQL database at {host}")
             return conn
     except Exception as e:
-        system.print_error(f"Failed to connect to PostgreSQL database at {host} with error: {e}")
+        system.print_error(f"Failed to connect to MySQL database at {host} with error: {e}")
 
-def check_data_patterns(conn, patterns, profile_name, database_name):
+def check_data_patterns(conn, patterns, profile_name, database_name, limit_start=0, limit_end=500, tables=None):
     cursor = conn.cursor()
-    cursor.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'")
+    
+    # Get the list of tables to scan
+    cursor.execute("SHOW TABLES")
     tables = [table[0] for table in cursor.fetchall()]
+    tables_to_scan = tables or []  # Use all tables if tables[] is blank or not provided
 
     table_count = 1
 
     results = []
-    for table in tables:
-        cursor.execute(f"SELECT * FROM {table}")
+    for table in tables_to_scan:
+        cursor.execute(f"SELECT * FROM {table} LIMIT {limit_end} OFFSET {limit_start}")
         columns = [column[0] for column in cursor.description]
 
         data_count = 1
@@ -42,15 +43,15 @@ def check_data_patterns(conn, patterns, profile_name, database_name):
                     if matches:
                         for match in matches:
                             results.append({
-                                'host': conn.dsn,
+                                'host': conn.get_host_info(),
                                 'database': database_name,
                                 'table': table,
                                 'column': column,
                                 'pattern_name': match['pattern_name'],
                                 'matches': match['matches'],
                                 'sample_text': match['sample_text'],
                                 'profile': profile_name,
-                                'data_source': 'postgresql'
+                                'data_source': 'mysql'
                             })
 
             data_count += 1
@@ -62,33 +63,40 @@ def check_data_patterns(conn, patterns, profile_name, database_name):
 
 def execute(args):
     results = []
-    system.print_info(f"Running Checks for PostgreSQL Sources")
+    system.print_info(f"Running Checks for MySQL Sources")
     connections = system.get_connection()
 
     if 'sources' in connections:
         sources_config = connections['sources']
-        postgresql_config = sources_config.get('postgresql')
+        mysql_config = sources_config.get('mysql')
 
-        if postgresql_config:
+        if mysql_config:
             patterns = system.get_fingerprint_file()
 
-            for key, config in postgresql_config.items():
+            for key, config in mysql_config.items():
                 host = config.get('host')
                 user = config.get('user')
-                port = config.get('port', 5432)  # default port for PostgreSQL
+                port = config.get('port', 3306)  # default port for MySQL
                 password = config.get('password')
                 database = config.get('database')
+                limit_start = config.get('limit_start', 0)
+                limit_end = config.get('limit_end', 500)
+                tables = config.get('tables', [])
 
                 if host and user and password and database:
-                    system.print_info(f"Checking PostgreSQL Profile {key} and database {database}")
-                    conn = connect_postgresql(host, port, user, password, database)
+                    system.print_info(f"Checking MySQL Profile {key} and database {database}")
+                    conn = connect_mysql(host, port, user, password, database)
                     if conn:
-                        results += check_data_patterns(conn, patterns, key, database)
+                        results += check_data_patterns(conn, patterns, key, database, limit_start=limit_start, limit_end=limit_end, tables=tables)
                         conn.close()
                 else:
-                    system.print_error(f"Incomplete PostgreSQL configuration for key: {key}")
+                    system.print_error(f"Incomplete MySQL configuration for key: {key}")
         else:
-            system.print_error("No PostgreSQL connection details found in connection.yml")
+            system.print_error("No MySQL connection details found in connection.yml")
     else:
         system.print_error("No 'sources' section found in connection.yml")
-    return results
+    return results
+
+# Example usage
+if __name__ == "__main__":
+    execute(None)
diff --git a/hawk_scanner/commands/postgresql.py b/hawk_scanner/commands/postgresql.py
@@ -1,7 +1,6 @@
 import psycopg2
 from hawk_scanner.internals import system
 from rich.console import Console
-from rich.table import Table
 
 console = Console()
 
@@ -20,14 +19,23 @@ def connect_postgresql(host, port, user, password, database):
     except Exception as e:
         system.print_error(f"Failed to connect to PostgreSQL database at {host} with error: {e}")
 
-def check_data_patterns(conn, patterns, profile_name, database_name):
+def check_data_patterns(conn, patterns, profile_name, database_name, limit_start=0, limit_end=500, tables=None):
     cursor = conn.cursor()
+    
+    # Get the list of tables to scan
     cursor.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'")
-    tables = [table[0] for table in cursor.fetchall()]
+    all_tables = [table[0] for table in cursor.fetchall()]
+    tables_to_scan = tables or all_tables  # Use all tables if tables[] is blank or not provided
+
     table_count = 1
+
     results = []
-    for table in tables:
-        cursor.execute(f"SELECT * FROM {table}")
+    for table in tables_to_scan:
+        if table not in all_tables:
+            system.print_warning(f"Table {table} not found in the database. Skipping.")
+            continue
+
+        cursor.execute(f"SELECT * FROM {table} LIMIT {limit_end} OFFSET {limit_start}")
         columns = [column[0] for column in cursor.description]
 
         data_count = 1
@@ -75,12 +83,15 @@ def execute(args):
                 port = config.get('port', 5432)  # default port for PostgreSQL
                 password = config.get('password')
                 database = config.get('database')
+                limit_start = config.get('limit_start', 0)
+                limit_end = config.get('limit_end', 500)
+                tables = config.get('tables', [])
 
                 if host and user and password and database:
-                    system.print_info(f"Checking PostgreSQL Profile {key} and database {database}")
+                    system.print_info(f"Checking PostgreSQL Profile {key}, database {database}")
                     conn = connect_postgresql(host, port, user, password, database)
                     if conn:
-                        results += check_data_patterns(conn, patterns, key, database)
+                        results += check_data_patterns(conn, patterns, key, database, limit_start=limit_start, limit_end=limit_end, tables=tables)
                         conn.close()
                 else:
                     system.print_error(f"Incomplete PostgreSQL configuration for key: {key}")
@@ -89,3 +100,7 @@ def execute(args):
     else:
         system.print_error("No 'sources' section found in connection.yml")
     return results
+
+# Example usage
+if __name__ == "__main__":
+    execute(None)
diff --git a/hawk_scanner/internals/system.py b/hawk_scanner/internals/system.py
@@ -1,3 +1,4 @@
+from androguard.misc import AnalyzeAPK
 from rich.console import Console 
 from rich.table import Table
 import json, requests, argparse, yaml, re, datetime, os, subprocess, platform, hashlib
@@ -266,6 +267,9 @@ def read_match_strings(file_path, source):
         # Check if the file is an archive (zip, rar, tar, tar.gz)
         elif file_path.lower().endswith(('.zip', '.rar', '.tar', '.tar.gz')):
             content = read_archive(file_path)
+        # Check if the file is an APK (Android application package)
+        elif file_path.lower().endswith('.apk'):
+            content = read_apk(file_path)
         else:
             # For other file types, read content normally
             with open(file_path, 'rb') as file:
@@ -278,6 +282,26 @@ def read_match_strings(file_path, source):
     matched_strings = match_strings(content)
     return matched_strings
 
+def read_apk(file_path):
+    try:
+        # Analyze the APK file using androguard
+        a, d, dx = AnalyzeAPK(file_path)
+
+        # Extract strings from the APK
+        strings = []
+        for method in dx.get_methods():
+            for _, _, _, string in method.get_strings():
+                strings.append(string)
+
+        # Combine all extracted strings into a single content string
+        content = ' '.join(strings)
+
+    except Exception as e:
+        print_debug(f"Error in read_apk: {e}")
+        content = ''
+
+    return content
+
 
 def read_pdf(file_path):
     content = ''
diff --git a/hawk_scanner/main.py b/hawk_scanner/main.py
diff --git a/readme.md b/readme.md
diff --git a/setup.py b/setup.py
