Pennyw0rth-NetExec/nxc/protocols/ssh/database.py

import configparser

from sqlalchemy import Boolean, Column, ForeignKeyConstraint, Integer, PrimaryKeyConstraint, String, select, func, delete
from sqlalchemy.dialects.sqlite import Insert
from sqlalchemy.orm import declarative_base

from nxc.database import BaseDB, format_host_query
from nxc.logger import nxc_logger
from nxc.paths import CONFIG_PATH

# we can't import config.py due to a circular dependency, so we have to create redundant code unfortunately
nxc_config = configparser.ConfigParser()
nxc_config.read(CONFIG_PATH)
nxc_workspace = nxc_config.get("nxc", "workspace", fallback="default")

Base = declarative_base()


class database(BaseDB):
    def __init__(self, db_engine):
        self.CredentialsTable = None
        self.HostsTable = None
        self.LoggedinRelationsTable = None
        self.AdminRelationsTable = None
        self.KeysTable = None

        super().__init__(db_engine)

    class Credential(Base):
        __tablename__ = "credentials"
        id = Column(Integer)
        username = Column(String)
        password = Column(String)
        credtype = Column(String)

        __table_args__ = (
            PrimaryKeyConstraint("id"),
        )

    class Host(Base):
        __tablename__ = "hosts"
        id = Column(Integer)
        host = Column(String)
        port = Column(Integer)
        banner = Column(String)
        os = Column(String)

        __table_args__ = (
            PrimaryKeyConstraint("id"),
        )

    class LoggedInRelation(Base):
        __tablename__ = "loggedin_relations"
        id = Column(Integer)
        credid = Column(Integer)
        hostid = Column(Integer)
        shell = Column(Boolean)

        __table_args__ = (
            PrimaryKeyConstraint("id"),
            ForeignKeyConstraint(["credid"], ["credentials.id"]),
            ForeignKeyConstraint(["hostid"], ["hosts.id"]),
        )

    # "admin" access with SSH means we have root access, which implies shell access since we run commands to check
    class AdminRelation(Base):
        __tablename__ = "admin_relations"
        id = Column(Integer)
        credid = Column(Integer)
        hostid = Column(Integer)

        __table_args__ = (
            PrimaryKeyConstraint("id"),
            ForeignKeyConstraint(["credid"], ["credentials.id"]),
            ForeignKeyConstraint(["hostid"], ["hosts.id"]),
        )

    class Key(Base):
        __tablename__ = "keys"
        id = Column(Integer)
        credid = Column(Integer)
        data = Column(String)

        __table_args__ = (
            PrimaryKeyConstraint("id"),
            ForeignKeyConstraint(["credid"], ["credentials.id"]),
        )

    @staticmethod
    def db_schema(db_conn):
        Base.metadata.create_all(db_conn)

    def reflect_tables(self):
        self.CredentialsTable = self.reflect_table(self.Credential)
        self.HostsTable = self.reflect_table(self.Host)
        self.LoggedinRelationsTable = self.reflect_table(self.LoggedInRelation)
        self.AdminRelationsTable = self.reflect_table(self.AdminRelation)
        self.KeysTable = self.reflect_table(self.Key)

    def add_host(self, host, port, banner, os=None):
        """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.host == host)
        results = self.db_execute(q).all()
        nxc_logger.debug(f"add_host(): Initial hosts results: {results}")

        # create new host
        if not results:
            new_host = {
                "host": host,
                "port": port,
                "banner": banner if banner is not None else "",
                "os": os if os is not None else "",
            }
            hosts = [new_host]
        # update existing hosts data
        else:
            for host_result in results:
                host_data = host_result._asdict()
                nxc_logger.debug(f"host: {host_result}")
                nxc_logger.debug(f"host_data: {host_data}")
                # only update column if it is being passed in
                if host is not None:
                    host_data["host"] = host
                if port is not None:
                    host_data["port"] = port
                if banner is not None:
                    host_data["banner"] = banner
                if os is not None:
                    host_data["os"] = os
                # 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"])
        nxc_logger.debug(f"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.db_execute(q, hosts)  # .scalar()
        # we only return updated IDs for now - when RETURNING clause is allowed we can return inserted
        if updated_ids:
            nxc_logger.debug(f"add_host() - Host IDs Updated: {updated_ids}")
            return updated_ids

    def add_credential(self, credtype, username, password, key=None):
        """Check if this credential has already been added to the database, if not add it in."""
        credentials = []

        # a user can have multiple keys, all with passphrases, and a separate login password
        if key is not None:
            q = (
                select(self.CredentialsTable)
                .join(self.KeysTable)
                .filter(
                    func.lower(self.CredentialsTable.c.username) == func.lower(username),
                    func.lower(self.CredentialsTable.c.credtype) == func.lower(credtype),
                    self.KeysTable.c.data == key,
                )
            )
            results = self.db_execute(q).all()
        else:
            q = select(self.CredentialsTable).filter(
                func.lower(self.CredentialsTable.c.username) == func.lower(username),
                func.lower(self.CredentialsTable.c.credtype) == func.lower(credtype),
            )
            results = self.db_execute(q).all()

        # add new credential
        if not results:
            new_cred = {
                "credtype": credtype,
                "username": username,
                "password": password,
            }
            credentials = [new_cred]
        # update existing cred data
        else:
            for creds in results:
                # this will include the id, so we don't touch it
                cred_data = creds._asdict()
                # only update column if it is being passed in
                if credtype is not None:
                    cred_data["credtype"] = credtype
                if username is not None:
                    cred_data["username"] = username
                if password is not None:
                    cred_data["password"] = password
                # only add cred to be updated if it has changed
                if cred_data not in credentials:
                    credentials.append(cred_data)

        # TODO: find a way to abstract this away to a single Upsert call
        q_users = Insert(self.CredentialsTable)  # .returning(self.CredentialsTable.c.id)
        update_columns_users = {col.name: col for col in q_users.excluded if col.name not in "id"}
        q_users = q_users.on_conflict_do_update(index_elements=self.CredentialsTable.primary_key,
                                                set_=update_columns_users)
        nxc_logger.debug(f"Adding credentials: {credentials}")

        self.db_execute(q_users, credentials)  # .scalar()

        # hacky way to get cred_id since we can't use returning() yet
        if len(credentials) == 1:
            cred_id = self.get_credential(credtype, username, password)
            if key is not None:
                self.add_key(cred_id, key)
            return cred_id
        else:
            return credentials

    def remove_credentials(self, creds_id):
        """Removes a credential ID from the database"""
        del_hosts = []
        for cred_id in creds_id:
            q = delete(self.CredentialsTable).filter(self.CredentialsTable.c.id == cred_id)
            del_hosts.append(q)
        self.db_execute(q)

    def add_key(self, cred_id, key):
        # check if key relation already exists
        check_q = self.db_execute(select(self.KeysTable).filter(self.KeysTable.c.credid == cred_id)).all()
        nxc_logger.debug(f"check_q: {check_q}")
        if check_q:
            nxc_logger.debug(f"Key already exists for cred_id {cred_id}")
            return None

        key_data = {"credid": cred_id, "data": key}
        self.db_execute(Insert(self.KeysTable), key_data)
        key_id = self.db_execute(select(self.KeysTable).filter(self.KeysTable.c.credid == cred_id)).all()[0].id
        nxc_logger.debug(f"Key added: {key_id}")
        return key_id

    def get_keys(self, key_id=None, cred_id=None):
        q = select(self.KeysTable)
        if key_id is not None:
            q = q.filter(self.KeysTable.c.id == key_id)
        elif cred_id is not None:
            q = q.filter(self.KeysTable.c.credid == cred_id)
        return self.db_execute(q).all()

    def add_admin_user(self, credtype, username, secret, host_id=None, cred_id=None):
        add_links = []

        creds_q = select(self.CredentialsTable)
        if cred_id:  # noqa: SIM108
            creds_q = creds_q.filter(self.CredentialsTable.c.id == cred_id)
        else:
            creds_q = creds_q.filter(
                func.lower(self.CredentialsTable.c.credtype) == func.lower(credtype),
                func.lower(self.CredentialsTable.c.username) == func.lower(username),
                self.CredentialsTable.c.password == secret,
            )
        creds = self.db_execute(creds_q)
        hosts = self.get_hosts(host_id)

        if creds and hosts:
            for cred, host in zip(creds, hosts, strict=True):
                cred_id = cred[0]
                host_id = host[0]
                link = {"credid": cred_id, "hostid": host_id}
                admin_relations_select = select(self.AdminRelationsTable).filter(
                    self.AdminRelationsTable.c.credid == cred_id,
                    self.AdminRelationsTable.c.hostid == host_id,
                )
                links = self.db_execute(admin_relations_select).all()

                if not links:
                    add_links.append(link)

        admin_relations_insert = Insert(self.AdminRelationsTable)

        if add_links:
            self.db_execute(admin_relations_insert, add_links)

    def get_admin_relations(self, cred_id=None, host_id=None):
        if cred_id:
            q = select(self.AdminRelationsTable).filter(self.AdminRelationsTable.c.credid == cred_id)
        elif host_id:
            q = select(self.AdminRelationsTable).filter(self.AdminRelationsTable.c.hostid == host_id)
        else:
            q = select(self.AdminRelationsTable)

        return self.db_execute(q).all()

    def remove_admin_relation(self, cred_ids=None, host_ids=None):
        q = delete(self.AdminRelationsTable)
        if cred_ids:
            for cred_id in cred_ids:
                q = q.filter(self.AdminRelationsTable.c.credid == cred_id)
        elif host_ids:
            for host_id in host_ids:
                q = q.filter(self.AdminRelationsTable.c.hostid == host_id)
        self.db_execute(q)

    def is_credential_valid(self, credential_id):
        """Check if this credential ID is valid."""
        q = select(self.CredentialsTable).filter(
            self.CredentialsTable.c.id == credential_id,
            self.CredentialsTable.c.password is not None,
        )
        results = self.db_execute(q).all()
        return len(results) > 0

    def get_credentials(self, filter_term=None, cred_type=None):
        """Return credentials from the database."""
        # if we're returning a single credential by ID
        if self.is_credential_valid(filter_term):
            q = select(self.CredentialsTable).filter(self.CredentialsTable.c.id == filter_term)
        elif cred_type:
            q = select(self.CredentialsTable).filter(self.CredentialsTable.c.credtype == cred_type)
        # if we're filtering by username
        elif filter_term and filter_term != "":
            like_term = func.lower(f"%{filter_term}%")
            q = select(self.CredentialsTable).filter(func.lower(self.CredentialsTable.c.username).like(like_term))
        # otherwise return all credentials
        else:
            q = select(self.CredentialsTable)

        return self.db_execute(q).all()

    def get_credential(self, cred_type, username, password):
        q = select(self.CredentialsTable).filter(
            self.CredentialsTable.c.username == username,
            self.CredentialsTable.c.password == password,
            self.CredentialsTable.c.credtype == cred_type,
        )
        results = self.db_execute(q).first()
        if results is not None:
            return results.id

    def is_host_valid(self, host_id):
        """Check if this host ID is valid."""
        q = select(self.HostsTable).filter(self.HostsTable.c.id == host_id)
        results = self.db_execute(q).all()
        return len(results) > 0

    def get_hosts(self, filter_term=None):
        """Return hosts from the database."""
        q = select(self.HostsTable)

        # if we're returning a single host by ID
        if self.is_host_valid(filter_term):
            q = q.filter(self.HostsTable.c.id == filter_term)
            results = self.db_execute(q).first()
            # all() returns a list, so we keep the return format the same so consumers don't have to guess
            return [results]
        # if we're filtering by host
        elif filter_term and filter_term != "":
            q = format_host_query(q, filter_term, self.HostsTable)

        results = self.db_execute(q).all()
        nxc_logger.debug(f"SSH get_hosts() - results: {results}")
        return results

    def is_user_valid(self, cred_id):
        """Check if this User ID is valid."""
        q = select(self.CredentialsTable).filter(self.CredentialsTable.c.id == cred_id)
        results = self.db_execute(q).all()
        return len(results) > 0

    def get_users(self, filter_term=None):
        q = select(self.CredentialsTable)

        if self.is_user_valid(filter_term):
            q = q.filter(self.CredentialsTable.c.id == filter_term)
        # if we're filtering by username
        elif filter_term and filter_term != "":
            like_term = func.lower(f"%{filter_term}%")
            q = q.filter(func.lower(self.CredentialsTable.c.username).like(like_term))
        return self.db_execute(q).all()

    def get_user(self, domain, username):
        q = select(self.CredentialsTable).filter(func.lower(self.CredentialsTable.c.username) == func.lower(username))
        return self.db_execute(q).all()

    def add_loggedin_relation(self, cred_id, host_id, shell=False):
        relation_query = select(self.LoggedinRelationsTable).filter(
            self.LoggedinRelationsTable.c.credid == cred_id,
            self.LoggedinRelationsTable.c.hostid == host_id,
        )
        results = self.db_execute(relation_query).all()

        # only add one if one doesn't already exist
        if not results:
            relation = {"credid": cred_id, "hostid": host_id, "shell": shell}
            try:
                nxc_logger.debug(f"Inserting loggedin_relations: {relation}")
                # TODO: find a way to abstract this away to a single Upsert call
                q = Insert(self.LoggedinRelationsTable)  # .returning(self.LoggedinRelationsTable.c.id)

                self.db_execute(q, [relation])  # .scalar()
                inserted_id_results = self.get_loggedin_relations(cred_id, host_id)
                nxc_logger.debug(f"Checking if relation was added: {inserted_id_results}")
                return inserted_id_results[0].id
            except Exception as e:
                nxc_logger.debug(f"Error inserting LoggedinRelation: {e}")

    def get_loggedin_relations(self, cred_id=None, host_id=None, shell=None):
        q = select(self.LoggedinRelationsTable)  # .returning(self.LoggedinRelationsTable.c.id)
        if cred_id:
            q = q.filter(self.LoggedinRelationsTable.c.credid == cred_id)
        if host_id:
            q = q.filter(self.LoggedinRelationsTable.c.hostid == host_id)
        if shell:
            q = q.filter(self.LoggedinRelationsTable.c.shell == shell)
        return self.db_execute(q).all()

    def remove_loggedin_relations(self, cred_id=None, host_id=None):
        q = delete(self.LoggedinRelationsTable)
        if cred_id:
            q = q.filter(self.LoggedinRelationsTable.c.credid == cred_id)
        elif host_id:
            q = q.filter(self.LoggedinRelationsTable.c.hostid == host_id)
        self.db_execute(q)