import pickle
import sqlite3


class SQLDatabase:
    """ `Database` management

        This class manages the `sqlite3 database`.
    """

    def __init__(self, db_abs_path, commands_table_name, jobs_table_name):
        """ Initialization of the `sqlite3 database`.

            This method initializes the `sqlite3` database creating 
            a table to store `allow` and `deny` OpenC2 commands and 
            a table to store non executed `APScheduler jobs` in case of an SLPF Actuator `shutdown`.

            :param db_abs_path: The absolute path to the database.
            :type db_abs_path: str
            :param commands_table_name: The `Commands` table name to store allow and deny OpenC2 commands.
            :type commands_table_name: str
            :param jobs_table_name: The `jobs` table name to store APScheduler jobs in case of an SLPF Actuator shutdown.
            :type jobs_table_name: str
        """
        try:
            self.db_abs_path = db_abs_path
            self.commands_table_name = commands_table_name
            self.jobs_table_name = jobs_table_name

            conn = sqlite3.connect(self.db_abs_path)
            c = conn.cursor()
            c.execute(f'''CREATE TABLE IF NOT EXISTS {self.commands_table_name} 
                    (rule_number INTEGER PRIMARY KEY,
                    action TEXT,
                    drop_process TEXT,
                    direction TEXT,
                    target TEXT,
                    protocol TEXT,
                    src_addr TEXT,
                    src_port INTEGER,
                    dst_addr TEXT,
                    dst_port INTEGER,
                    start_time TEXT,
                    stop_time TEXT,
                    persistent INTEGER,
                    custom_data BLOB,
                    scheduler_data BLOB)''')

            c.execute(f'''CREATE TABLE IF NOT EXISTS {self.jobs_table_name}
                    (id TEXT PRIMARY KEY,
                    func_name TEXT,
                    next_run_time TEXT,
                    args BLOB,
                    kwargs BLOB)''')

            conn.commit()
        except sqlite3.Error as e:
            raise e
        finally:
            conn.close()

    def insert_command(self, insert_rule, action, drop_process, direction, target, protocol, src_addr, src_port,
                       dst_addr, dst_port, start_time, stop_time, persistent, custom_data, scheduler_data):
        """ This method stores allow and deny OpenC2 `commands` into the database `commands table`.

            The `insert_rule` parameter could have a `None` value, in this case the rule number of the command is automaticaly generated by the database.

            :param insert_rule: The rule number associated with the command.
            :type insert_rule: int
            :param action: The name of the command action to be stored in the database.
            :type action: str
            :param drop_process: Specifies how to handle denied packets.
            :type drop_process: str
            :param direction: The name of the command direction to be stored in the database.
            :type direction: str
            :param target: The name of the command target to be stored in the database.
            :type target: str
            :param protocol: The name of the command protocol to be stored in the database.
            :type protocol: str
            :param src_addr: The command IPv4/IPv6 source address to be stored in the database.
            :type src_addr: str
            :param src_port: The command source port to be stored in the database.
            :type src_port: int
            :param dst_addr: The command IPv4/IPv6 destination address to be stored in the database.
            :type dst_addr: str
            :param dst_port: The command destination port to be stored in the database.
            :type dst_port: int
            :param start_time: The time at which the command is activated.
            :type start_time: str
            :param stop_time: The time at which the command is deactivated.
            :type stop_time: str
            :param persistent: Specifies if the command is persistent or not.
            :type persistent: bool
            :param custom_data: Actuator-specific custom data associated with the rule to be saved.
            :type custom_data: dict
            :param scheduler_data: Contains information about `APScheduler job IDs`: 
                                    `start_job_id` is the ID of the job responsible to activate, at a certain `start time`, the OpenC2 allow or deny action, 
                                    `stop_job_id` is the ID of the job responsible to deactivate, at a certain `stop time`, the OpenC2 allow or deny action. 
            :type job_ids: dict

            :return: The rule number associated with the stored command.
            :rtype: int
        """

        try:
            conn = sqlite3.connect(self.db_abs_path)
            custom_data_blob = pickle.dumps(custom_data)
            scheduler_data_blob = pickle.dumps(scheduler_data)
            c = conn.cursor()
            c.execute(f'''INSERT INTO {self.commands_table_name} 
                          (rule_number, action, drop_process, direction, target, protocol, src_addr, src_port, dst_addr, dst_port, start_time, stop_time, persistent, custom_data, scheduler_data) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)''',
                      (insert_rule, action, drop_process, direction, target, protocol, src_addr, src_port, dst_addr,
                       dst_port, start_time, stop_time, persistent, custom_data_blob, scheduler_data_blob))
            rule_number = c.lastrowid
            conn.commit()
            return rule_number
        except sqlite3.Error as e:
            raise e
        finally:
            conn.close()

    def get_command(self, rule_number):
        """ This method retrieves the command with the given rule number from the database `commands table`.

            The returned command is a `Python dictionary`, it has to be transformed into an `OpenC2 Command`.

            :param rule_number: The rule number associated with the command to retrieve.
            :type rule_number: int

            :return: The needed command.
            :rtype: dict
        """
        try:
            conn = sqlite3.connect(self.db_abs_path)
            c = conn.cursor()
            c.execute(f'SELECT * FROM {self.commands_table_name} WHERE rule_number = ?', (rule_number,))
            db_result = c.fetchone()
            return self.recreate_command(db_result)
        except sqlite3.Error as e:
            raise e
        finally:
            conn.close()

    def delete_command(self, rule_number):
        """ This method deletes the command with the given rule number from the database `commands table`.

            :param rule_number: The rule number associated with the command to delete.
            :type rule_number: int
        """
        try:
            conn = sqlite3.connect(self.db_abs_path)
            c = conn.cursor()
            c.execute(f'DELETE FROM {self.commands_table_name} WHERE rule_number = ?', (rule_number,))
            conn.commit()
        except sqlite3.Error as e:
            raise e
        finally:
            conn.close()

    def find_command(self, rule_number):
        """ This method checks whether the given rule number is already associated with a command in the database `commands table`.

            :param rule_number: The rule number associated with the command to delete.
            :type rule_number: int

            :return: This method returns `True` if the given rule number is already associated with a command, `False` otherwise.
            :rtype: bool
        """
        try:
            conn = sqlite3.connect(self.db_abs_path)
            c = conn.cursor()
            c.execute(f'SELECT EXISTS(SELECT 1 FROM {self.commands_table_name} WHERE rule_number = ?)', (rule_number,))
            found = c.fetchone()[0]
            return found

        except sqlite3.Error as e:
            raise e
        finally:
            conn.close()

    def get_non_persistent_comands(self):
        """ This method gets all `non-persistent commands` from the database `commands table`.

            The returned list of non-persistent commands is a list of `Python dictionaries`, it has to be transformed into a list of `OpenC2 Commands`.

            :return: The list of non-persistent commands.
            :rtype: list
        """
        try:
            conn = sqlite3.connect(self.db_abs_path)
            c = conn.cursor()
            c.execute(f'SELECT * FROM {self.commands_table_name} WHERE persistent = 0')
            non_persistent = c.fetchall()
            result = []
            for command in non_persistent:
                result.append(self.recreate_command(command))
            return result
        except sqlite3.Error as e:
            raise e
        finally:
            conn.close()

    def recreate_command(self, command):
        """ This method creates a `Python dictonary` representing the command from `database data`.

            :param command: The database data representing the command.
            :type command: tuple

            :return: The `Python dictonary` representing the command.
            :rtype: dict
        """
        return {
            'rule_number': command[0],
            'action': command[1],
            'drop_process': command[2],
            'direction': command[3],
            'target': command[4],
            'protocol': command[5],
            'src_addr': command[6],
            'src_port': command[7],
            'dst_addr': command[8],
            'dst_port': command[9],
            'start_time': command[10],
            'stop_time': command[11],
            'persistent': command[12],
            'custom_data': pickle.loads(command[13]),
            'scheduler_data': pickle.loads(command[14])
        }

    def insert_job(self, id, func_name, next_run_time, args, kwargs):
        """ This method stores an `APScheduler job` into the database `jobs table`.

            :param id: The id of the APScheduler job.
            :type id: str
            :param func_name: The name of the function that the APScheduler job has to execute.
            :type func_name: str
            :param next_run_time: The time at which the job must execute the function.
            :type next_run_time: str
            :param args: A list of arguments to pass to the function to be executed.
            :type next_run_time: list
            :param kwargs: A dictionary of arguments to pass to the function to be executed.
            :type kwargs: dict
        """
        try:
            conn = sqlite3.connect(self.db_abs_path)

            args_blob = pickle.dumps(args)
            kwargs_blob = pickle.dumps(kwargs)

            c = conn.cursor()
            c.execute(f'''INSERT INTO {self.jobs_table_name} 
                          (id, func_name, next_run_time, args, kwargs) VALUES (?, ?, ?, ?, ?)''',
                      (id, func_name, next_run_time, args_blob, kwargs_blob))
            conn.commit()
        except sqlite3.Error as e:
            raise e
        finally:
            conn.close()

    def get_jobs(self):
        """ This method gets all `APScheduler jobs` from the database `jobs table`.

            :return: A list of Python dictionaries representing the APScheduler jobs.
            :rtype: list
        """
        try:
            conn = sqlite3.connect(self.db_abs_path)
            c = conn.cursor()
            c.execute(f'SELECT * FROM {self.jobs_table_name}')
            jobs = c.fetchall()
            result = []
            for job in jobs:
                result.append(self.recreate_job(job))
            return result
        except sqlite3.Error as e:
            raise e
        finally:
            conn.close()

    def delete_jobs(self):
        """ This method deletes all `APScheduler jobs` from the database `jobs table`. """
        try:
            conn = sqlite3.connect(self.db_abs_path)
            c = conn.cursor()
            c.execute(f'DELETE FROM {self.jobs_table_name}')
            conn.commit()
        except sqlite3.Error as e:
            raise e
        finally:
            conn.close()

    def recreate_job(self, job):
        """ This method creates a `Python dictonary` representing the APScheduler job from `database data`.

            :param command: The database data representing the job.
            :type command: tuple

            :return: The `Python dictonary` representing the job.
            :rtype: dict
        """
        return {
            'id': job[0],
            'func_name': job[1],
            'next_run_time': job[2],
            'args': pickle.loads(job[3]),
            'kwargs': pickle.loads(job[4])
        }

    def clean_db(self):
        """ This method deletes all `OpenC2 Commands` from the database `commands table` and all `APScheduler jobs` from the database `jobs table`. """
        try:
            conn = sqlite3.connect(self.db_abs_path)
            c = conn.cursor()
            c.execute(f'DELETE FROM {self.commands_table_name}')
            c.execute(f'DELETE FROM {self.jobs_table_name}')
            conn.commit()
        except sqlite3.Error as e:
            raise e
        finally:
            conn.close()

    def is_empty(self):
        """ This method checks whether the database `commands tables` is empty. 
        
            :return: `True` if the database is empty, `False` otherwise.
            :rtype: bool 
        """
        try:
            conn = sqlite3.connect(self.db_abs_path)
            c = conn.cursor()
            c.execute(f'SELECT 1 FROM {self.commands_table_name} LIMIT 1')
            found = c.fetchone()
            return found is None
        except sqlite3.Error as e:
            raise e
        finally:
            conn.close()
