Merge branch 'main' into fix-dbs

Author: NeffIsBack

nxc/connection.py

@@ -206,13 +206,16 @@ def print_host_info(self):
     def create_conn_obj(self):
         return
 
+    def disconnect(self):
+        return
+
     def check_if_admin(self):
         return
 
     def kerberos_login(self, domain, username, password="", ntlm_hash="", aesKey="", kdcHost="", useCache=False):
         return
 
-    def plaintext_login(self, domain, username, password):
+    def plaintext_login(self, username, password):
         return
 
     def hash_login(self, domain, username, ntlm_hash):
@@ -234,6 +237,7 @@ def proto_flow(self):
                 else:
                     self.logger.debug("Calling command arguments")
                     self.call_cmd_args()
+            self.disconnect()
 
     def call_cmd_args(self):
         """Calls all the methods specified by the command line arguments

nxc/database.py

@@ -51,24 +51,8 @@ def write_configfile(config, config_path):
         config.write(configfile)
 
 
-def create_workspace(workspace_name, p_loader=None):
-    """
-    Create a new workspace with the given name.
-
-    Args:
-    ----
-        workspace_name (str): The name of the workspace.
-
-    Returns:
-    -------
-        None
-    """
-    if exists(path_join(WORKSPACE_DIR, workspace_name)):
-        print(f"[-] Workspace {workspace_name} already exists")
-    else:
-        print(f"[*] Creating {workspace_name} workspace")
-        mkdir(path_join(WORKSPACE_DIR, workspace_name))
-
+def init_protocol_dbs(workspace_name, p_loader=None):
+    """Check for each protocol if the database exists, if not create it."""
     if p_loader is None:
         p_loader = ProtocolLoader()
     protocols = p_loader.get_protocols()
@@ -93,6 +77,27 @@ def create_workspace(workspace_name, p_loader=None):
             conn.close()
 
 
+def create_workspace(workspace_name, p_loader=None):
+    """
+    Create a new workspace with the given name.
+
+    Args:
+    ----
+        workspace_name (str): The name of the workspace.
+
+    Returns:
+    -------
+        None
+    """
+    if exists(path_join(WORKSPACE_DIR, workspace_name)):
+        print(f"[-] Workspace {workspace_name} already exists")
+    else:
+        print(f"[*] Creating {workspace_name} workspace")
+        mkdir(path_join(WORKSPACE_DIR, workspace_name))
+
+    init_protocol_dbs(workspace_name, p_loader)
+
+
 def delete_workspace(workspace_name):
     shutil.rmtree(path_join(WORKSPACE_DIR, workspace_name))
     print(f"[*] Workspace {workspace_name} deleted")
@@ -137,3 +142,7 @@ def db_execute(self, *args):
         res = self.sess.execute(*args)
         self.lock.release()
         return res
+
+
+    # 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")