fix: check if an IP is being searched for when calling get_hosts in db calls; fixes Pennyw0rth#589

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,19 @@ 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
+    try:
+        ipaddress.ip_address(filter_term)
+        nxc_logger.debug(f"filter_term is an IP address: {filter_term}")
+        q = q.filter(HostsTable.c.ip == filter_term)
+    except ValueError:
+        nxc_logger.debug(f"filter_term is not an IP address: {filter_term}")
+        like_term = func.lower(f"%{filter_term}%")
+        q = q.filter(HostsTable.c.ip.like(like_term) | func.lower(HostsTable.c.hostname).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