Merge pull request #65 from idnahacks/dev

Allow the sql-path arg to be a dir or file

Author: idnahacks

goodhound/__init__.py

@@ -25,7 +25,7 @@ def arguments():
     parsegroupschema.add_argument("--patch41", help="A temporary option to patch a bug in Bloodhound 4.1 relating to the highvalue attribute.", action="store_true")
     parsegroupsql = argparser.add_argument_group('SQLite Database')
     parsegroupsql.add_argument("--db-skip", help="Skips the logging of attack paths to a local SQLite Database", action="store_true")
-    parsegroupsql.add_argument("-sqlpath", "--sql-path", default=getcwd(), help="Sets the file path of the SQLite Database file goodhound.db", type=str)
+    parsegroupsql.add_argument("-sqlpath", "--sql-path", default=getcwd(), help="Sets the file path of the GoodHound Database file", type=str)
     args = argparser.parse_args()
     return args
 

goodhound/ghutils.py

@@ -34,4 +34,16 @@ def checkoutdir(path):
             logging.error('Selected output path is not a directory')
             exit(1)
     else:
-        os.makedirs(path, exist_ok=True)
+        os.makedirs(path, exist_ok=True)
+
+def checkdbfileexists(sqlpath):
+    """Looks at the provided sql-path argument and determines whether to create a new db or update an existing one."""
+    if Path(sqlpath).exists():
+        if Path(sqlpath).is_file():
+            dbfile = sqlpath
+        else:
+            dbfile = sqlpath + os.sep + "goodhound.db"
+    else:
+        os.makedirs(sqlpath, exist_ok=True)
+        dbfile = sqlpath + os.sep + "goodhound.db"
+    return dbfile

goodhound/sqldb.py

@@ -1,4 +1,4 @@
-from goodhound import neodb
+from goodhound import neodb, ghutils
 from sqlite3.dbapi2 import Error
 import sqlite3
 from pathlib import Path
@@ -20,9 +20,9 @@ def db(results, graph, args):
         first_seen INTEGER NOT NULL,
     	last_seen INTEGER NOT NULL);"""
         conn = None
+        dbpath = ghutils.checkdbfileexists(args.sql_path)
         try:
-            db = str(Path(args.sql_path)) + os.sep + 'goodhound.db'
-            conn = sqlite3.connect(db)
+            conn = sqlite3.connect(dbpath)
             c = conn.cursor()
             c.execute(table_sql)
             scandate, scandatenice = neodb.getscandate(graph)