restructuring libs to be a bit more sane in terms of schema and so that we can split between aws and gcp

libs/aws/sql.py

@@ -0,0 +1,128 @@
+'''
+Custom SQL/database functions for WeirdAAL
+'''
+
+import sqlite3
+from sqlite3 import Error
+
+
+def create_table(db_name, table_name, sql):
+    '''
+    SQLite3 create table function
+    '''
+    with sqlite3.connect(db_name) as db:
+        cursor = db.cursor()
+        cursor.execute("""SELECT name FROM sqlite_master WHERE name=?""", (table_name,))
+        result = cursor.fetchall()
+        keep_table = True
+        if len(result) == 1:
+            response = input("The table {} already exists, do you wish to recreate it? (y/n): ".format(table_name))
+            if response == "y":
+                keep_table = False
+                print("The {} table will be recreated - all existing data will be lost".format(table_name))
+                cursor.execute("drop table if exists {}".format(table_name))
+                db.commit()
+            else:
+                print("The existing table was kept")
+        else:
+            keep_table = False
+        if not keep_table:
+            cursor.execute(sql)
+            db.commit()
+
+
+def create_recon_table(db_name, table_name):
+    '''
+    Create recon table service:subservice:AWSKeyID,time
+    '''
+    sql = """CREATE TABLE recon
+             (ID integer,
+             service text,
+             sub_service text,
+             AWSKeyID text,
+             target text,
+             checked_at timestamp,
+             PRIMARY KEY (ID))"""
+             # FOREIGN KEY (AWSKeyID) references AWSKey(ID))"""
+    create_table(db_name, table_name, sql)
+    print("created table: {}".format(table_name))
+
+
+def create_awskey_table(db_name, table_name):
+    '''
+    Create awskey table (currently unused)
+    '''
+    sql = """CREATE TABLE AWSKey
+             (ID integer,
+             AWSKeyID text,
+             description text,
+             target text,
+             PRIMARY KEY(ID))"""
+    create_table(db_name, table_name, sql)
+    print("created table: {}".format(table_name))
+
+
+def create_services_table(db_name, table_name):
+    '''
+    Create services table - service:sub_service:sub_service_data
+    '''
+    sql = """CREATE TABLE services
+             (ID integer,
+             AWSKeyID Text,
+             service text,
+             sub_service text,
+             sub_service_data text,
+             checked_at timestamp,
+             target text,
+             PRIMARY KEY(ID))"""
+    create_table(db_name, table_name, sql)
+    print("created table: {}".format(table_name))
+
+
+def insert_awskey_data(db_name, records):
+    '''
+    Insert AWS Key and a description to the AWSKey table (unused)
+    '''
+    sql = """INSERT INTO AWSKey(AWSKeyID, description, target) VALUES (?,?,?)"""
+    for record in records:
+        query(db_name, sql, record)
+
+
+def insert_reconservice_data(db_name, records):
+    '''
+    Insert data into the recon table
+    '''
+    sql = """INSERT INTO recon(service, sub_service, AWSKeyID, target, checked_at) VALUES (?,?,?,?,?)"""
+    for record in records:
+        query(db_name, sql, record)
+
+
+def insert_sub_service_data(db_name, records):
+    '''
+    Insert service, sub_service & sub_service data into the DB
+    '''
+    sql = """INSERT INTO services(service, sub_service, sub_service_data, AWSKeyID, target, checked_at) VALUES (?,?,?,?,?,?)"""
+    for record in records:
+        query(db_name, sql, record)
+
+
+def search_recon_by_key(db_name, AWSKeyID):
+    '''
+    Function to query services by AWSKey and order them by time
+    '''
+    with sqlite3.connect(db_name) as db:
+        cursor = db.cursor()
+        cursor.execute("""SELECT DISTINCT service, sub_service, checked_at FROM recon WHERE AWSKeyID=? ORDER BY datetime(checked_at)""", (AWSKeyID,))
+        results = cursor.fetchall()
+        return results
+
+
+def query(db_name, sql, data):
+    '''
+    Generic query function
+    '''
+    with sqlite3.connect(db_name) as db:
+        cursor = db.cursor()
+        # cursor.execute("""PRAGMA foreign_keys = ON""")
+        cursor.execute(sql, data)
+        db.commit()