aswi2021jmsd-gitlab/src/dao/certificate_repository.py @ 3243b35d

import time
from sqlite3 import Connection, Error, DatabaseError, IntegrityError, ProgrammingError, OperationalError, \
    NotSupportedError
from typing import Dict, List, Set

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
        """

        Logger.debug("Function launched.")

        try:
            sql = (f"INSERT INTO {TAB_CERTIFICATES} "
                   f"({COL_VALID_FROM},"
                   f"{COL_VALID_TO},"
                   f"{COL_PEM_DATA},"
                   f"{COL_COMMON_NAME},"
                   f"{COL_COUNTRY_CODE},"
                   f"{COL_LOCALITY},"
                   f"{COL_PROVINCE},"
                   f"{COL_ORGANIZATION},"
                   f"{COL_ORGANIZATIONAL_UNIT},"
                   f"{COL_EMAIL_ADDRESS},"
                   f"{COL_TYPE_ID},"
                   f"{COL_PARENT_ID},"
                   f"{COL_PRIVATE_KEY_ID}) "
                   f"VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?)")
            values = [certificate.valid_from,
                      certificate.valid_to,
                      certificate.pem_data,
                      certificate.common_name,
                      certificate.country_code,
                      certificate.locality,
                      certificate.province,
                      certificate.organization,
                      certificate.organizational_unit,
                      certificate.email_address,
                      certificate.type_id,
                      certificate.parent_id,
                      certificate.private_key_id]
            self.cursor.execute(sql, values)
            self.connection.commit()

            last_id: int = self.cursor.lastrowid

            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 (IntegrityError, OperationalError, ProgrammingError, NotSupportedError, DatabaseError, Error) as e:
            Logger.error(str(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
        """

        Logger.debug("Function launched.")

        try:
            sql = (f"SELECT * FROM {TAB_CERTIFICATES} "
                   f"WHERE {COL_ID} = ? AND ({COL_DELETION_DATE} IS NULL OR {COL_DELETION_DATE} = '')")
            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],      # ID
                                                   certificate_row[1],      # valid from
                                                   certificate_row[2],      # valid to
                                                   certificate_row[3],      # pem data
                                                   certificate_row[14],     # type ID
                                                   certificate_row[15],     # parent ID
                                                   certificate_row[16],     # private key ID
                                                   usage_dict,
                                                   certificate_row[4],      # common name
                                                   certificate_row[5],      # country code
                                                   certificate_row[6],      # locality
                                                   certificate_row[7],      # province
                                                   certificate_row[8],      # organization
                                                   certificate_row[9],      # organizational unit
                                                   certificate_row[10],     # email address
                                                   certificate_row[11],     # revocation date
                                                   certificate_row[12])     # revocation reason

        except (IntegrityError, OperationalError, ProgrammingError, NotSupportedError, DatabaseError, Error) as e:
            Logger.error(str(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
        """

        Logger.debug("Function launched.")

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

            sql = (f"SELECT * FROM {TAB_CERTIFICATES} "
                   f"WHERE ({COL_DELETION_DATE} IS NULL OR {COL_DELETION_DATE} = ''){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],      # ID
                                                certificate_row[1],      # valid from
                                                certificate_row[2],      # valid to
                                                certificate_row[3],      # pem data
                                                certificate_row[14],     # type ID
                                                certificate_row[15],     # parent ID
                                                certificate_row[16],     # private key ID
                                                usage_dict,
                                                certificate_row[4],      # common name
                                                certificate_row[5],      # country code
                                                certificate_row[6],      # locality
                                                certificate_row[7],      # province
                                                certificate_row[8],      # organization
                                                certificate_row[9],      # organizational unit
                                                certificate_row[10],     # email address
                                                certificate_row[11],     # revocation date
                                                certificate_row[12]))    # revocation reason

        except (IntegrityError, OperationalError, ProgrammingError, NotSupportedError, DatabaseError, Error) as e:
            Logger.error(str(e))
            raise DatabaseException(e)

        return certificates

    def read_all_filter(self, target_types: Set[int], target_usages: Set[int], target_cn_substring: str, page: int,
                        per_page: int):
        """
        Reads (selects) all certificates according to a specific filtering and pagination options.
        :param target_types: certificate types (filter)
        :param target_usages: certificate usages (filter)
        :param target_cn_substring: certificate CN substring (filter)
        :param page: target page
        :param per_page: target page size
        :return: list of certificates
        """

        Logger.debug("Function launched.")

        try:
            values = []
            values += list(target_types)

            sql = (
                f"SELECT * "
                f"FROM {TAB_CERTIFICATES} a "
                f"WHERE (a.{COL_DELETION_DATE} IS NULL OR a.{COL_DELETION_DATE} = '') "
                f"AND a.{COL_TYPE_ID} IN ({','.join('?' * len(target_types))}) "
            )
            
            if target_usages is not None:
                sql += f"AND (SELECT COUNT(*) FROM {TAB_CERTIFICATE_USAGES} WHERE {COL_CERTIFICATE_ID} = a.{COL_ID} " \
                    f"AND {COL_USAGE_TYPE_ID} IN ({','.join('?' * len(target_usages))})) > 0 "
                values += list(target_usages)

            if target_cn_substring is not None:
                sql += f" AND a.{COL_COMMON_NAME} LIKE '%' || ? || '%'"

                values += [target_cn_substring]

            if page is not None and per_page is not None:
                sql += f" LIMIT {per_page} OFFSET {page * per_page}"

            self.cursor.execute(sql, values)
            certificate_rows = self.cursor.fetchall()

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

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

                certificates.append(Certificate(certificate_row[0],  # ID
                                                certificate_row[1],  # valid from
                                                certificate_row[2],  # valid to
                                                certificate_row[3],  # pem data
                                                certificate_row[14],  # type ID
                                                certificate_row[15],  # parent ID
                                                certificate_row[16],  # private key ID
                                                usage_dict,
                                                certificate_row[4],  # common name
                                                certificate_row[5],  # country code
                                                certificate_row[6],  # locality
                                                certificate_row[7],  # province
                                                certificate_row[8],  # organization
                                                certificate_row[9],  # organizational unit
                                                certificate_row[10],  # email address
                                                certificate_row[11],  # revocation date
                                                certificate_row[12]))  # revocation reason

        except (IntegrityError, OperationalError, ProgrammingError, NotSupportedError, DatabaseError, Error) as e:
            Logger.error(str(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
        """

        Logger.debug("Function launched.")

        try:
            sql = (f"UPDATE {TAB_CERTIFICATES} "
                   f"SET {COL_VALID_FROM} = ?, "
                   f"{COL_VALID_TO} = ?, "
                   f"{COL_PEM_DATA} = ?, "
                   f"{COL_COMMON_NAME} = ?, "
                   f"{COL_COUNTRY_CODE} = ?, "
                   f"{COL_LOCALITY} = ?, "
                   f"{COL_PROVINCE} = ?, "
                   f"{COL_ORGANIZATION} = ?, "
                   f"{COL_ORGANIZATIONAL_UNIT} = ?, "
                   f"{COL_EMAIL_ADDRESS} = ?, "
                   f"{COL_TYPE_ID} = ?, "
                   f"{COL_PARENT_ID} = ?, "
                   f"{COL_PRIVATE_KEY_ID} = ? "
                   f"WHERE {COL_ID} = ?")
            values = [certificate.valid_from,
                      certificate.valid_to,
                      certificate.pem_data,
                      certificate.common_name,
                      certificate.country_code,
                      certificate.locality,
                      certificate.province,
                      certificate.organization,
                      certificate.organizational_unit,
                      certificate.email_address,
                      certificate.type_id,
                      certificate.parent_id,
                      certificate.private_key_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()

            check_updated = self.cursor.rowcount > 0

            # 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 (IntegrityError, OperationalError, ProgrammingError, NotSupportedError, DatabaseError, Error) as e:
            Logger.error(str(e))
            raise DatabaseException(e)

        return check_updated

    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
        """

        Logger.debug("Function launched.")

        try:
            sql = (f"UPDATE {TAB_CERTIFICATES} "
                   f"SET {COL_DELETION_DATE} = ? "
                   f"WHERE {COL_ID} = ? AND ({COL_DELETION_DATE} IS NULL OR {COL_DELETION_DATE} = '')")
            values = [int(time.time()),
                      certificate_id]
            self.cursor.execute(sql, values)
            self.connection.commit()

        except (IntegrityError, OperationalError, ProgrammingError, NotSupportedError, DatabaseError, Error) as e:
            Logger.error(str(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
        """

        Logger.debug("Function launched.")

        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 (IntegrityError, OperationalError, ProgrammingError, NotSupportedError, DatabaseError, Error) as e:
            Logger.error(str(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
        """

        Logger.debug("Function launched.")

        try:
            sql = (f"UPDATE {TAB_CERTIFICATES} "
                   f"SET {COL_REVOCATION_DATE} = NULL, "
                   f"{COL_REVOCATION_REASON} = NULL "
                   f"WHERE {COL_ID} = ?")
            values = [certificate_id]
            self.cursor.execute(sql, values)
            self.connection.commit()

        except (IntegrityError, OperationalError, ProgrammingError, NotSupportedError, DatabaseError, Error) as e:
            Logger.error(str(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
        """

        Logger.debug("Function launched.")

        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],      # ID
                                                certificate_row[1],      # valid from
                                                certificate_row[2],      # valid to
                                                certificate_row[3],      # pem data
                                                certificate_row[14],     # type ID
                                                certificate_row[15],     # parent ID
                                                certificate_row[16],     # private key ID
                                                usage_dict,
                                                certificate_row[4],      # common name
                                                certificate_row[5],      # country code
                                                certificate_row[6],      # locality
                                                certificate_row[7],      # province
                                                certificate_row[8],      # organization
                                                certificate_row[9],      # organizational unit
                                                certificate_row[10],     # email address
                                                certificate_row[11],     # revocation date
                                                certificate_row[12]))    # revocation reason

        except (IntegrityError, OperationalError, ProgrammingError, NotSupportedError, DatabaseError, Error) as e:
            Logger.error(str(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
        """

        Logger.debug("Function launched.")

        try:
            sql = (f"SELECT * FROM {TAB_CERTIFICATES} "
                   f"WHERE {COL_PARENT_ID} = ? AND {COL_ID} != ? AND "
                   f"({COL_DELETION_DATE} IS NULL OR {COL_DELETION_DATE} = '')")
            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],      # ID
                                                certificate_row[1],      # valid from
                                                certificate_row[2],      # valid to
                                                certificate_row[3],      # pem data
                                                certificate_row[14],     # type ID
                                                certificate_row[15],     # parent ID
                                                certificate_row[16],     # private key ID
                                                usage_dict,
                                                certificate_row[4],      # common name
                                                certificate_row[5],      # country code
                                                certificate_row[6],      # locality
                                                certificate_row[7],      # province
                                                certificate_row[8],      # organization
                                                certificate_row[9],      # organizational unit
                                                certificate_row[10],     # email address
                                                certificate_row[11],     # revocation date
                                                certificate_row[12]))    # revocation reason

        except (IntegrityError, OperationalError, ProgrammingError, NotSupportedError, DatabaseError, Error) as e:
            Logger.error(str(e))
            raise DatabaseException(e)

        return certificates

    def get_all_issued_by_filter(self, issuer_id, target_types, target_usages, target_cn_substring, page, per_page):
        """
        Get list of the certificates that are direct descendants of the certificate with the ID

        :param target_types: certificate types (filter)
        :param target_usages: certificate usages (filter)
        :param target_cn_substring: certificate CN substring (filter)
        :param page: target page
        :param per_page: target page size
        :param issuer_id: ID of specific certificate

        :return:
            list of the certificates OR
            None if the list is empty
        """

        Logger.debug("Function launched.")

        try:
            values = [issuer_id, issuer_id]
            values += list(target_types)

            sql = (
                f"SELECT * "
                f"FROM {TAB_CERTIFICATES} a "
                f"WHERE (a.{COL_DELETION_DATE} IS NULL OR a.{COL_DELETION_DATE} = '') "
                f"AND a.{COL_PARENT_ID} = ? AND a.{COL_ID} != ? "
                f"AND a.{COL_TYPE_ID} IN ({','.join('?' * len(target_types))}) "
                
            )
            
            if target_usages is not None:
                sql += f"AND (SELECT COUNT(*) FROM {TAB_CERTIFICATE_USAGES} WHERE {COL_CERTIFICATE_ID} = a.{COL_ID} " \
                    f"AND {COL_USAGE_TYPE_ID} IN ({','.join('?' * len(target_usages))})) > 0 "
                values += list(target_usages)

            if target_cn_substring is not None:
                sql += f" AND a.{COL_COMMON_NAME} LIKE '%' || ? || '%'"
                values += [target_cn_substring]

            if page is not None and per_page is not None:
                sql += f" LIMIT {per_page} OFFSET {page * per_page}"

            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],  # ID
                                                certificate_row[1],  # valid from
                                                certificate_row[2],  # valid to
                                                certificate_row[3],  # pem data
                                                certificate_row[14],  # type ID
                                                certificate_row[15],  # parent ID
                                                certificate_row[16],  # private key ID
                                                usage_dict,
                                                certificate_row[4],  # common name
                                                certificate_row[5],  # country code
                                                certificate_row[6],  # locality
                                                certificate_row[7],  # province
                                                certificate_row[8],  # organization
                                                certificate_row[9],  # organizational unit
                                                certificate_row[10],  # email address
                                                certificate_row[11],  # revocation date
                                                certificate_row[12]))  # revocation reason

        except (IntegrityError, OperationalError, ProgrammingError, NotSupportedError, DatabaseError, Error) as e:
            Logger.error(str(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
        """

        Logger.debug("Function launched.")

        try:
            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)

        except (IntegrityError, OperationalError, ProgrammingError, NotSupportedError, DatabaseError, Error) as e:
            Logger.error(str(e))
            raise DatabaseException(e)

        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
        """

        Logger.debug("Function launched.")

        try:
            # 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

        except (IntegrityError, OperationalError, ProgrammingError, NotSupportedError, DatabaseError, Error) as e:
            Logger.error(str(e))
            raise DatabaseException(e)

        return 1