aswi2021jmsd-gitlab/src/dao/certificate_repository.py @ 5e31b492

from sqlite3 import Connection, Error
from typing import Dict, List

from src.exceptions.database_exception import DatabaseException
from injector import inject
from src.constants import *
from src.model.certificate import Certificate
from src.utils.logger import Logger


class CertificateRepository:

    @inject
    def __init__(self, connection: Connection):
        """
        Constructor of the CertificateRepository object

        :param connection: Instance of the Connection object
        """
        self.connection = connection
        self.cursor = connection.cursor()

    def create(self, certificate: Certificate):
        """
        Creates a certificate.
        For root certificate (CA) the parent certificate id is modified to the same id (id == parent_id).

        :param certificate: Instance of the Certificate object

        :return: the result of whether the creation was successful
        """

        try:
            sql = (f"INSERT INTO {TAB_CERTIFICATES} "
                   f"({COL_COMMON_NAME},"
                   f"{COL_VALID_FROM},"
                   f"{COL_VALID_TO},"
                   f"{COL_PEM_DATA},"
                   f"{COL_PRIVATE_KEY_ID},"
                   f"{COL_TYPE_ID},"
                   f"{COL_PARENT_ID})"
                   f"VALUES(?,?,?,?,?,?,?)")
            values = [certificate.common_name,
                      certificate.valid_from,
                      certificate.valid_to,
                      certificate.pem_data,
                      certificate.private_key_id,
                      certificate.type_id,
                      certificate.parent_id]
            self.cursor.execute(sql, values)
            self.connection.commit()

            last_id: int = self.cursor.lastrowid

            # TODO assure that this is correct
            if certificate.type_id == ROOT_CA_ID:
                certificate.parent_id = last_id
                self.update(last_id, certificate)
            else:
                for usage_id, usage_value in certificate.usages.items():
                    if usage_value:
                        sql = (f"INSERT INTO {TAB_CERTIFICATE_USAGES} "
                               f"({COL_CERTIFICATE_ID},"
                               f"{COL_USAGE_TYPE_ID}) "
                               f"VALUES (?,?)")
                        values = [last_id, usage_id]
                        self.cursor.execute(sql, values)
                        self.connection.commit()
        except Error as e:
            raise DatabaseException(e)

        return last_id

    def read(self, certificate_id: int):
        """
        Reads (selects) a certificate.

        :param certificate_id: ID of specific certificate

        :return: instance of the Certificate object
        """

        try:
            sql = (f"SELECT * FROM {TAB_CERTIFICATES} "
                   f"WHERE {COL_ID} = ?")
            values = [certificate_id]
            self.cursor.execute(sql, values)
            certificate_row = self.cursor.fetchone()

            if certificate_row is None:
                return None

            sql = (f"SELECT * FROM {TAB_CERTIFICATE_USAGES} "
                   f"WHERE {COL_CERTIFICATE_ID} = ?")
            self.cursor.execute(sql, values)
            usage_rows = self.cursor.fetchall()

            usage_dict: Dict[int, bool] = {}
            for usage_row in usage_rows:
                usage_dict[usage_row[2]] = True

            certificate: Certificate = Certificate(certificate_row[0],
                                                   certificate_row[1],
                                                   certificate_row[2],
                                                   certificate_row[3],
                                                   certificate_row[4],
                                                   certificate_row[7],
                                                   certificate_row[8],
                                                   certificate_row[9],
                                                   usage_dict,
                                                   certificate_row[5],
                                                   certificate_row[6])
        except Error as e:
            raise DatabaseException(e)

        return certificate

    def read_all(self, filter_type: int = None):
        """
        Reads (selects) all certificates (with type).

        :param filter_type: ID of certificate type from CertificateTypes table

        :return: list of certificates
        """

        try:
            sql_extension = ""
            values = []
            if filter_type is not None:
                sql_extension = (f" WHERE {COL_TYPE_ID} = ("
                                 f"SELECT {COL_ID} FROM {TAB_CERTIFICATE_TYPES} WHERE {COL_ID} = ?)")
                values = [filter_type]

            sql = f"SELECT * FROM {TAB_CERTIFICATES}{sql_extension}"
            self.cursor.execute(sql, values)
            certificate_rows = self.cursor.fetchall()

            certificates: List[Certificate] = []
            for certificate_row in certificate_rows:
                sql = (f"SELECT * FROM {TAB_CERTIFICATE_USAGES} "
                       f"WHERE {COL_CERTIFICATE_ID} = ?")
                values = [certificate_row[0]]
                self.cursor.execute(sql, values)
                usage_rows = self.cursor.fetchall()

                usage_dict: Dict[int, bool] = {}
                for usage_row in usage_rows:
                    usage_dict[usage_row[2]] = True

                certificates.append(Certificate(certificate_row[0],
                                                certificate_row[1],
                                                certificate_row[2],
                                                certificate_row[3],
                                                certificate_row[4],
                                                certificate_row[7],
                                                certificate_row[8],
                                                certificate_row[9],
                                                usage_dict,
                                                certificate_row[5],
                                                certificate_row[6]))
        except Error as e:
            raise DatabaseException(e)

        return certificates

    def update(self, certificate_id: int, certificate: Certificate):
        """
        Updates a certificate.
        If the parameter of certificate (Certificate object) is not to be changed,
        the same value must be specified.

        :param certificate_id: ID of specific certificate
        :param certificate: Instance of the Certificate object

        :return: the result of whether the updation was successful
        """

        try:
            sql = (f"UPDATE {TAB_CERTIFICATES} "
                   f"SET {COL_COMMON_NAME} = ?, "
                   f"{COL_VALID_FROM} = ?, "
                   f"{COL_VALID_TO} = ?, "
                   f"{COL_PEM_DATA} = ?, "
                   f"{COL_PRIVATE_KEY_ID} = ?, "
                   f"{COL_TYPE_ID} = ?, "
                   f"{COL_PARENT_ID} = ? "
                   f"WHERE {COL_ID} = ?")
            values = [certificate.common_name,
                      certificate.valid_from,
                      certificate.valid_to,
                      certificate.pem_data,
                      certificate.private_key_id,
                      certificate.type_id,
                      certificate.parent_id,
                      certificate_id]
            self.cursor.execute(sql, values)
            self.connection.commit()

            sql = (f"DELETE FROM {TAB_CERTIFICATE_USAGES} "
                   f"WHERE {COL_CERTIFICATE_ID} = ?")
            values = [certificate_id]
            self.cursor.execute(sql, values)
            self.connection.commit()

            # iterate over usage pairs
            for usage_id, usage_value in certificate.usages.items():
                if usage_value:
                    sql = (f"INSERT INTO {TAB_CERTIFICATE_USAGES} "
                           f"({COL_CERTIFICATE_ID},"
                           f"{COL_USAGE_TYPE_ID}) "
                           f"VALUES (?,?)")
                    values = [certificate_id, usage_id]
                    self.cursor.execute(sql, values)
                    self.connection.commit()
        except Error as e:
            raise DatabaseException(e)

        return self.cursor.rowcount > 0

    def delete(self, certificate_id: int):
        """
        Deletes a certificate

        :param certificate_id: ID of specific certificate

        :return: the result of whether the deletion was successful
        """

        try:
            sql = (f"DELETE FROM {TAB_CERTIFICATES} "
                   f"WHERE {COL_ID} = ?")
            values = [certificate_id]
            self.cursor.execute(sql, values)
            self.connection.commit()
        except Error as e:
            raise DatabaseException(e)

        return self.cursor.rowcount > 0

    def set_certificate_revoked(
            self, certificate_id: int, revocation_date: str, revocation_reason: str = REV_REASON_UNSPECIFIED):
        """
        Revoke a certificate

        :param certificate_id: ID of specific certificate
        :param revocation_date: Date, when the certificate is revoked
        :param revocation_reason: Reason of the revocation

        :return:
            the result of whether the revocation was successful OR
            sqlite3.Error if an exception is thrown
        """

        try:
            if revocation_date != "" and revocation_reason == "":
                revocation_reason = REV_REASON_UNSPECIFIED
            elif revocation_date == "":
                return False

            sql = (f"UPDATE {TAB_CERTIFICATES} "
                   f"SET {COL_REVOCATION_DATE} = ?, "
                   f"{COL_REVOCATION_REASON} = ? "
                   f"WHERE {COL_ID} = ? AND ({COL_REVOCATION_DATE} IS NULL OR {COL_REVOCATION_DATE} = '')")
            values = [revocation_date,
                      revocation_reason,
                      certificate_id]
            self.cursor.execute(sql, values)
            self.connection.commit()
        except Error as e:
            raise DatabaseException(e)

        return self.cursor.rowcount > 0

    def clear_certificate_revocation(self, certificate_id: int):
        """
        Clear revocation of a certificate

        :param certificate_id: ID of specific certificate

        :return:
            the result of whether the clear revocation was successful OR
            sqlite3.Error if an exception is thrown
        """

        try:
            sql = (f"UPDATE {TAB_CERTIFICATES} "
                   f"SET {COL_REVOCATION_DATE} = '', "
                   f"{COL_REVOCATION_REASON} = '' "
                   f"WHERE {COL_ID} = ?")
            values = [certificate_id]
            self.cursor.execute(sql, values)
            self.connection.commit()
        except Error as e:
            raise DatabaseException(e)

        return self.cursor.rowcount > 0

    def get_all_revoked_by(self, certificate_id: int):
        """
        Get list of the revoked certificates that are direct descendants of the certificate with the ID

        :param certificate_id: ID of specific certificate

        :return:
            list of the certificates OR
            None if the list is empty OR
            sqlite3.Error if an exception is thrown
        """

        try:
            sql = (f"SELECT * FROM {TAB_CERTIFICATES} "
                   f"WHERE {COL_PARENT_ID} = ? AND {COL_REVOCATION_DATE} IS NOT NULL AND {COL_REVOCATION_DATE} != ''")
            values = [certificate_id]
            self.cursor.execute(sql, values)
            certificate_rows = self.cursor.fetchall()

            certificates: List[Certificate] = []
            for certificate_row in certificate_rows:
                sql = (f"SELECT * FROM {TAB_CERTIFICATE_USAGES} "
                       f"WHERE {COL_CERTIFICATE_ID} = ?")
                values = [certificate_row[0]]
                self.cursor.execute(sql, values)
                usage_rows = self.cursor.fetchall()

                usage_dict: Dict[int, bool] = {}
                for usage_row in usage_rows:
                    usage_dict[usage_row[2]] = True

                certificates.append(Certificate(certificate_row[0],
                                                certificate_row[1],
                                                certificate_row[2],
                                                certificate_row[3],
                                                certificate_row[4],
                                                certificate_row[7],
                                                certificate_row[8],
                                                certificate_row[9],
                                                usage_dict,
                                                certificate_row[5],
                                                certificate_row[6]))
        except Error as e:
            raise DatabaseException(e)

        return certificates

    def get_all_issued_by(self, certificate_id: int):
        """
        Get list of the certificates that are direct descendants of the certificate with the ID

        :param certificate_id: ID of specific certificate

        :return:
            list of the certificates OR
            None if the list is empty OR
            sqlite3.Error if an exception is thrown
        """

        try:
            sql = (f"SELECT * FROM {TAB_CERTIFICATES} "
                   f"WHERE {COL_PARENT_ID} = ? AND {COL_ID} != ?")
            values = [certificate_id, certificate_id]
            self.cursor.execute(sql, values)
            certificate_rows = self.cursor.fetchall()

            certificates: List[Certificate] = []
            for certificate_row in certificate_rows:
                sql = (f"SELECT * FROM {TAB_CERTIFICATE_USAGES} "
                       f"WHERE {COL_CERTIFICATE_ID} = ?")
                values = [certificate_row[0]]
                self.cursor.execute(sql, values)
                usage_rows = self.cursor.fetchall()

                usage_dict: Dict[int, bool] = {}
                for usage_row in usage_rows:
                    usage_dict[usage_row[2]] = True

                certificates.append(Certificate(certificate_row[0],
                                                certificate_row[1],
                                                certificate_row[2],
                                                certificate_row[3],
                                                certificate_row[4],
                                                certificate_row[7],
                                                certificate_row[8],
                                                certificate_row[9],
                                                usage_dict,
                                                certificate_row[5],
                                                certificate_row[6]))
        except Error as e:
            raise DatabaseException(e)

        return certificates

    def get_all_descendants_of(self, certificate_id: int):
        """
        Get a list of all certificates C such that the certificate identified by "certificate_id" belongs to a trust chain
        between C and its root certificate authority (i.e. is an ancestor of C).
        :param certificate_id: target certificate ID
        :return: list of all descendants
        """
        def dfs(children_of, this, collection: list):
            for child in children_of(this.certificate_id):
                dfs(children_of, child, collection)
            collection.append(this)

        subtree_root = self.read(certificate_id)
        if subtree_root is None:
            return None

        all_certs = []
        dfs(self.get_all_issued_by, subtree_root, all_certs)
        return all_certs

    def get_next_id(self) -> int:
        """
        Get identifier of the next certificate that will be inserted into the database
        :return: identifier of the next certificate that will be added into the database
        """
        # get next IDs of all tables
        self.cursor.execute("SELECT * FROM SQLITE_SEQUENCE")
        results = self.cursor.fetchall()

        # search for next ID in Certificates table and return it
        for result in results:
            if result[0] == TAB_CERTIFICATES:
                return result[1] + 1  # current last id + 1
        # if certificates table is not present in the query results, return 1
        return 1