feat(ssh): initial db update for SSH to properly add in hosts
Marshall Hallenbeck
parent
3798682ed7
commit
3a7f0305e0
|
|
@ -1,10 +1,20 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
from sqlalchemy.dialects.sqlite import Insert
|
||||
from sqlalchemy.orm import sessionmaker, scoped_session
|
||||
from sqlalchemy import MetaData, Table
|
||||
from sqlalchemy import MetaData, Table, select
|
||||
from sqlalchemy.exc import IllegalStateChangeError, NoInspectionAvailable, NoSuchTableError
|
||||
|
||||
import os
|
||||
import configparser
|
||||
|
||||
from cme.logger import cme_logger
|
||||
from cme.paths import CME_PATH
|
||||
|
||||
# we can't import config.py due to a circular dependency, so we have to create redundant code unfortunately
|
||||
cme_config = configparser.ConfigParser()
|
||||
cme_config.read(os.path.join(CME_PATH, "cme.conf"))
|
||||
cme_workspace = cme_config.get("CME", "workspace", fallback="default")
|
||||
|
||||
|
||||
class database:
|
||||
|
|
@ -21,41 +31,66 @@ class database:
|
|||
)
|
||||
|
||||
Session = scoped_session(session_factory)
|
||||
# this is still named "conn" when it is the session object; TODO: rename
|
||||
self.conn = Session()
|
||||
self.sess = Session()
|
||||
|
||||
@staticmethod
|
||||
def db_schema(db_conn):
|
||||
db_conn.execute('''CREATE TABLE "credentials" (
|
||||
db_conn.execute(
|
||||
'''CREATE TABLE "credentials" (
|
||||
"id" integer PRIMARY KEY,
|
||||
"username" text,
|
||||
"password" text
|
||||
)''')
|
||||
|
||||
db_conn.execute('''CREATE TABLE "hosts" (
|
||||
"password" text,
|
||||
"credtype" text
|
||||
)''')
|
||||
db_conn.execute(
|
||||
'''CREATE TABLE "hosts" (
|
||||
"id" integer PRIMARY KEY,
|
||||
"ip" text,
|
||||
"hostname" text,
|
||||
"port" integer,
|
||||
"server_banner" text
|
||||
)''')
|
||||
)''')
|
||||
db_conn.execute(
|
||||
'''CREATE TABLE "loggedin_relations" (
|
||||
"id" integer PRIMARY KEY,
|
||||
"userid" integer,
|
||||
"hostid" integer,
|
||||
FOREIGN KEY(userid) REFERENCES users(id),
|
||||
FOREIGN KEY(hostid) REFERENCES hosts(id)
|
||||
)''')
|
||||
db_conn.execute(
|
||||
'''CREATE TABLE "admin_relations" (
|
||||
"id" integer PRIMARY KEY,
|
||||
"userid" integer,
|
||||
"hostid" integer,
|
||||
FOREIGN KEY(userid) REFERENCES users(id),
|
||||
FOREIGN KEY(hostid) REFERENCES hosts(id)
|
||||
)''')
|
||||
db_conn.execute(
|
||||
'''CREATE TABLE "keys" (
|
||||
"id" integer PRIMARY KEY,
|
||||
"data" text,
|
||||
FOREIGN KEY(userid) REFERENCES users(id)
|
||||
)''')
|
||||
|
||||
def reflect_tables(self):
|
||||
with self.db_engine.connect() as conn:
|
||||
with self.db_engine.connect():
|
||||
try:
|
||||
self.CredentialsTable = Table("credentials", self.metadata, autoload_with=self.db_engine)
|
||||
self.HostsTable = Table("hosts", self.metadata, autoload_with=self.db_engine)
|
||||
except (NoInspectionAvailable, NoSuchTableError):
|
||||
ssh_workspace = f"~/.cme/workspaces/{cme_workspace}/ssh.db"
|
||||
print(
|
||||
"[-] Error reflecting tables - this means there is a DB schema mismatch \n"
|
||||
"[-] Error reflecting tables for SSH protocol - this means there is a DB schema mismatch \n"
|
||||
"[-] This is probably because a newer version of CME is being ran on an old DB schema\n"
|
||||
"[-] If you wish to save the old DB data, copy it to a new location (`cp -r ~/.cme/workspaces/ ~/old_cme_workspaces/`)\n"
|
||||
"[-] Then remove the CME DB folders (`rm -rf ~/.cme/workspaces/`) and rerun CME to initialize the new DB schema"
|
||||
f"[-] Optionally save the old DB data (`cp {ssh_workspace} ~/cme_ssh.bak`)\n"
|
||||
f"[-] Then remove the CME SSH DB (`rm -rf {ssh_workspace}`) and run CME to initialize the new DB"
|
||||
)
|
||||
exit()
|
||||
|
||||
def shutdown_db(self):
|
||||
try:
|
||||
self.conn.close()
|
||||
self.sess.close()
|
||||
# due to the async nature of CME, sometimes session state is a bit messy and this will throw:
|
||||
# Method 'close()' can't be called here; method '_connection_for_bind()' is already in progress and
|
||||
# this would cause an unexpected state change to <SessionTransactionState.CLOSED: 5>
|
||||
|
|
@ -64,4 +99,64 @@ class database:
|
|||
|
||||
def clear_database(self):
|
||||
for table in self.metadata.sorted_tables:
|
||||
self.conn.execute(table.delete())
|
||||
self.sess.execute(table.delete())
|
||||
|
||||
def add_host(self, ip, hostname, port, os, banner):
|
||||
"""
|
||||
Check if this host has already been added to the database, if not, add it in.
|
||||
"""
|
||||
hosts = []
|
||||
updated_ids = []
|
||||
|
||||
q = select(self.HostsTable).filter(
|
||||
self.HostsTable.c.ip == ip
|
||||
)
|
||||
results = self.conn.execute(q).all()
|
||||
|
||||
# create new host
|
||||
if not results:
|
||||
new_host = {
|
||||
"ip": ip,
|
||||
"hostname": hostname if hostname is not None else '',
|
||||
"port": port,
|
||||
"os": os if os is not None else '',
|
||||
"banner": banner if banner is not None else ''
|
||||
}
|
||||
hosts = [new_host]
|
||||
# update existing hosts data
|
||||
else:
|
||||
for host in results:
|
||||
host_data = host._asdict()
|
||||
# only update column if it is being passed in
|
||||
if ip is not None:
|
||||
host_data["ip"] = ip
|
||||
if hostname is not None:
|
||||
host_data["hostname"] = hostname
|
||||
if port is not None:
|
||||
host_data["port"] = port
|
||||
if os is not None:
|
||||
host_data["os"] = os
|
||||
if banner is not None:
|
||||
host_data["banner"] = banner
|
||||
# only add host to be updated if it has changed
|
||||
if host_data not in hosts:
|
||||
hosts.append(host_data)
|
||||
updated_ids.append(host_data["id"])
|
||||
cme_logger.debug(f"Update Hosts: {hosts}")
|
||||
|
||||
# TODO: find a way to abstract this away to a single Upsert call
|
||||
q = Insert(self.HostsTable) # .returning(self.HostsTable.c.id)
|
||||
update_columns = {col.name: col for col in q.excluded if col.name not in 'id'}
|
||||
q = q.on_conflict_do_update(
|
||||
index_elements=self.HostsTable.primary_key,
|
||||
set_=update_columns
|
||||
)
|
||||
|
||||
self.conn.execute(
|
||||
q,
|
||||
hosts
|
||||
) # .scalar()
|
||||
# we only return updated IDs for now - when RETURNING clause is allowed we can return inserted
|
||||
if updated_ids:
|
||||
cme_logger.debug(f"add_host() - Host IDs Updated: {updated_ids}")
|
||||
return updated_ids
|
||||
|
|
|
|||
Loading…
Reference in New Issue