NetExec/cme/protocols/mssql/database.py

193 lines
6.3 KiB
Python
Executable File

class database:
def __init__(self, conn):
self.conn = conn
@staticmethod
def db_schema(db_conn):
db_conn.execute('''CREATE TABLE "computers" (
"id" integer PRIMARY KEY,
"ip" text,
"hostname" text,
"domain" text,
"os" text,
"instances" integer
)''')
# This table keeps track of which credential has admin access over which machine and vice-versa
db_conn.execute('''CREATE TABLE "admin_relations" (
"id" integer PRIMARY KEY,
"userid" integer,
"computerid" integer,
FOREIGN KEY(userid) REFERENCES users(id),
FOREIGN KEY(computerid) REFERENCES computers(id)
)''')
# type = hash, plaintext
db_conn.execute('''CREATE TABLE "users" (
"id" integer PRIMARY KEY,
"credtype" text,
"domain" text,
"username" text,
"password" text
)''')
def add_computer(self, ip, hostname, domain, os, instances):
"""
Check if this host has already been added to the database, if not add it in.
"""
cur = self.conn.cursor()
cur.execute('SELECT * FROM computers WHERE ip LIKE ?', [ip])
results = cur.fetchall()
if not len(results):
cur.execute("INSERT INTO computers (ip, hostname, domain, os, instances) VALUES (?,?,?,?,?)", [ip, hostname, domain, os, instances])
cur.close()
def add_credential(self, credtype, domain, username, password):
"""
Check if this credential has already been added to the database, if not add it in.
"""
cur = self.conn.cursor()
cur.execute("SELECT * FROM users WHERE credtype=? AND LOWER(domain)=LOWER(?) AND LOWER(username)=LOWER(?) AND password=?", [credtype, domain, username, password])
results = cur.fetchall()
if not len(results):
cur.execute("INSERT INTO users (credtype, domain, username, password) VALUES (?,?,?,?)", [credtype, domain, username, password])
cur.close()
def remove_credentials(self, credIDs):
"""
Removes a credential ID from the database
"""
for credID in credIDs:
cur = self.conn.cursor()
cur.execute("DELETE FROM users WHERE id=?", [credID])
cur.close()
def add_admin_user(self, credtype, domain, username, password, host):
cur = self.conn.cursor()
cur.execute("SELECT * FROM users WHERE credtype=? AND LOWER(domain)=LOWER(?) AND LOWER(username)=LOWER(?) AND password=?", [credtype, domain, username, password])
creds = cur.fetchall()
cur.execute('SELECT * FROM computers WHERE ip LIKE ?', [host])
hosts = cur.fetchall()
if len(creds) and len(hosts):
for cred, host in zip(creds, hosts):
userid = cred[0]
computerid = host[0]
# Check to see if we already added this link
cur.execute("SELECT * FROM admin_relations WHERE userid=? AND computerid=?", [userid, computerid])
links = cur.fetchall()
if not len(links):
cur.execute("INSERT INTO admin_relations (userid, computerid) VALUES (?,?)", [userid, computerid])
cur.close()
def get_admin_relations(self, userID=None, hostID=None):
cur = self.conn.cursor()
if userID:
cur.execute("SELECT * from admin_relations WHERE userid=?", [userID])
elif hostID:
cur.execute("SELECT * from admin_relations WHERE computerid=?", [hostID])
results = cur.fetchall()
cur.close()
return results
def remove_admin_relation(self, userIDs=None, hostIDs=None):
cur = self.conn.cursor()
if userIDs:
for userID in userIDs:
cur.execute("DELETE FROM admin_relations WHERE userid=?", [userID])
elif hostIDs:
for hostID in hostIDs:
cur.execute("DELETE FROM admin_relations WHERE computerid=?", [hostID])
cur.close()
def is_credential_valid(self, credentialID):
"""
Check if this credential ID is valid.
"""
cur = self.conn.cursor()
cur.execute('SELECT * FROM users WHERE id=? LIMIT 1', [credentialID])
results = cur.fetchall()
cur.close()
return len(results) > 0
def get_credentials(self, filterTerm=None, credtype=None):
"""
Return credentials from the database.
"""
cur = self.conn.cursor()
# if we're returning a single credential by ID
if self.is_credential_valid(filterTerm):
cur.execute("SELECT * FROM users WHERE id=? LIMIT 1", [filterTerm])
# if we're filtering by credtype
elif credtype:
cur.execute("SELECT * FROM users WHERE credtype=?", [credtype])
# if we're filtering by username
elif filterTerm and filterTerm != "":
cur.execute("SELECT * FROM users WHERE LOWER(username) LIKE LOWER(?)", ['%{}%'.format(filterTerm.lower())])
# otherwise return all credentials
else:
cur.execute("SELECT * FROM users")
results = cur.fetchall()
cur.close()
return results
def is_computer_valid(self, hostID):
"""
Check if this computer ID is valid.
"""
cur = self.conn.cursor()
cur.execute('SELECT * FROM computers WHERE id=? LIMIT 1', [hostID])
results = cur.fetchall()
cur.close()
return len(results) > 0
def get_computers(self, filterTerm=None):
"""
Return computers from the database.
"""
cur = self.conn.cursor()
# if we're returning a single host by ID
if self.is_computer_valid(filterTerm):
cur.execute("SELECT * FROM computers WHERE id=? LIMIT 1", [filterTerm])
# if we're filtering by ip/hostname
elif filterTerm and filterTerm != "":
cur.execute("SELECT * FROM computers WHERE ip LIKE ? OR LOWER(hostname) LIKE LOWER(?)", ['%{}%'.format(filterTerm.lower()), '%{}%'.format(filterTerm.lower())])
# otherwise return all credentials
else:
cur.execute("SELECT * FROM computers")
results = cur.fetchall()
cur.close()
return results