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