Merge pull request Pennyw0rth#590 from Pennyw0rth/marshall-db-ip-fix

fix: check if an IP is being searched for when calling get_hosts

Author: Marshall-Hallenbeck

nxc/database.py

@@ -1,4 +1,5 @@
 import configparser
+import ipaddress
 import shutil
 import sys
 from os import mkdir
@@ -8,7 +9,7 @@
 from sqlite3 import connect
 from threading import Lock
 
-from sqlalchemy import create_engine, MetaData
+from sqlalchemy import create_engine, MetaData, func
 from sqlalchemy.exc import IllegalStateChangeError
 from sqlalchemy.orm import sessionmaker, scoped_session
 
@@ -109,7 +110,36 @@ def initialize_db():
 
     # Even if the default workspace exists, we still need to check if every protocol has a database (in case of a new protocol)
     init_protocol_dbs("default")
-
+    
+def format_host_query(q, filter_term, HostsTable):
+    """One annoying thing is that if you search for an ip such as '10.10.10.5',
+    it will return 10.10.10.5 and 10.10.10.52, so we have to check if its an ip address first
+    """
+    # the FTP and SSH protocols call the column host instead of IP
+    # TODO: normalize these column names
+    if hasattr(HostsTable.c, "ip"):
+        ip_column = HostsTable.c.ip
+        nxc_logger.debug("Using 'ip' column for filtering")
+    elif hasattr(HostsTable.c, "host"):
+        ip_column = HostsTable.c.host
+        nxc_logger.debug("Using 'host' column for filtering")
+    else:
+        nxc_logger.debug("Neither 'ip' nor 'host' columns found in the table")
+        return q
+   
+    # first we check if its an ip address 
+    try:
+        ipaddress.ip_address(filter_term)
+        nxc_logger.debug(f"filter_term is an IP address: {filter_term}")
+        q = q.filter(ip_column == filter_term)
+    except ValueError:
+        nxc_logger.debug(f"filter_term is not an IP address: {filter_term}")
+        like_term = func.lower(f"%{filter_term}%")
+        
+        # check if the hostname column exists for hostname searching
+        q = q.filter(ip_column.like(like_term) | func.lower(HostsTable.c.hostname).like(like_term)) if hasattr(HostsTable.c, "hostname") else q.filter(ip_column.like(like_term))
+
+    return q
 
 class BaseDB:
     def __init__(self, db_engine):

nxc/protocols/ftp/database.py

@@ -7,7 +7,7 @@
     NoSuchTableError,
 )
 
-from nxc.database import BaseDB
+from nxc.database import BaseDB, format_host_query
 from nxc.logger import nxc_logger
 
 
@@ -221,8 +221,8 @@ def get_hosts(self, filter_term=None):
             return [results]
         # if we're filtering by host
         elif filter_term and filter_term != "":
-            like_term = func.lower(f"%{filter_term}%")
-            q = q.filter(self.HostsTable.c.host.like(like_term))
+            q = format_host_query(q, filter_term, self.HostsTable)
+
         results = self.db_execute(q).all()
         nxc_logger.debug(f"FTP get_hosts() - results: {results}")
         return results

nxc/protocols/mssql/database.py

@@ -5,7 +5,7 @@
 from sqlalchemy.dialects.sqlite import Insert  # used for upsert
 from sqlalchemy.exc import SAWarning, NoInspectionAvailable, NoSuchTableError
 
-from nxc.database import BaseDB
+from nxc.database import BaseDB, format_host_query
 from nxc.logger import nxc_logger
 
 # if there is an issue with SQLAlchemy and a connection cannot be cleaned up properly it spews out annoying warnings
@@ -272,7 +272,6 @@ def get_hosts(self, filter_term=None, domain=None):
                 q = q.filter(func.lower(self.HostsTable.c.domain) == func.lower(domain))
         # if we're filtering by ip/hostname
         elif filter_term and filter_term != "":
-            like_term = func.lower(f"%{filter_term}%")
-            q = select(self.HostsTable).filter(self.HostsTable.c.ip.like(like_term) | func.lower(self.HostsTable.c.hostname).like(like_term))
+            q = format_host_query(q, filter_term, self.HostsTable)
 
         return self.db_execute(q).all()

nxc/protocols/smb/database.py

@@ -11,7 +11,7 @@
 )
 from sqlalchemy.exc import SAWarning
 
-from nxc.database import BaseDB
+from nxc.database import BaseDB, format_host_query
 from nxc.logger import nxc_logger
 
 # if there is an issue with SQLAlchemy and a connection cannot be cleaned up properly it spews out annoying warnings
@@ -349,6 +349,7 @@ def add_admin_user(self, credtype, domain, username, password, host, user_id=Non
         hosts = self.get_hosts(host)
 
         if users and hosts:
+            nxc_logger.debug(f"users: {users}, hosts: {hosts}")
             for user, host in zip(users, hosts, strict=True):
                 user_id = user[0]
                 host_id = host[0]
@@ -468,8 +469,8 @@ def get_hosts(self, filter_term=None, domain=None):
             q = q.filter(self.HostsTable.c.domain.like(like_term))
         # if we're filtering by ip/hostname
         elif filter_term and filter_term != "":
-            like_term = func.lower(f"%{filter_term}%")
-            q = q.filter(self.HostsTable.c.ip.like(like_term) | func.lower(self.HostsTable.c.hostname).like(like_term))
+            q = format_host_query(q, filter_term, self.HostsTable)
+
         results = self.db_execute(q).all()
         nxc_logger.debug(f"smb hosts() - results: {results}")
         return results

nxc/protocols/ssh/database.py

@@ -9,7 +9,7 @@
     NoSuchTableError,
 )
 
-from nxc.database import BaseDB
+from nxc.database import BaseDB, format_host_query
 from nxc.logger import nxc_logger
 from nxc.paths import NXC_PATH
 
@@ -348,8 +348,8 @@ def get_hosts(self, filter_term=None):
             return [results]
         # if we're filtering by host
         elif filter_term and filter_term != "":
-            like_term = func.lower(f"%{filter_term}%")
-            q = q.filter(self.HostsTable.c.host.like(like_term))
+            q = format_host_query(q, filter_term, self.HostsTable)
+
         results = self.db_execute(q).all()
         nxc_logger.debug(f"SSH get_hosts() - results: {results}")
         return results

nxc/protocols/winrm/database.py

@@ -1,13 +1,12 @@
 import sys
-
 from sqlalchemy import Table, select, func, delete
 from sqlalchemy.dialects.sqlite import Insert
 from sqlalchemy.exc import (
     NoInspectionAvailable,
     NoSuchTableError,
 )
 
-from nxc.database import BaseDB
+from nxc.database import BaseDB, format_host_query
 from nxc.logger import nxc_logger
 
 
@@ -309,8 +308,8 @@ def get_hosts(self, filter_term=None):
             q = q.filter(self.HostsTable.c.domain.like(like_term))
         # if we're filtering by ip/hostname
         elif filter_term and filter_term != "":
-            like_term = func.lower(f"%{filter_term}%")
-            q = q.filter(self.HostsTable.c.ip.like(like_term) | func.lower(self.HostsTable.c.hostname).like(like_term))
+            q = format_host_query(q, filter_term, self.HostsTable)
+
         results = self.db_execute(q).all()
         nxc_logger.debug(f"winrm get_hosts() - results: {results}")
         return results