reportmonster/classes/vessel.py

from classes.database import Database

from configparser import SectionProxy
from typing import Optional, Union
from datetime import datetime

from MySQLdb.cursors import DictCursor

from const import *


class Vessel:
    """Class describing a Vessel
    """
    @classmethod
    def fromConfig(cls, config: SectionProxy):
        """Create Vessel object from a Vessel section in the Config file

        Args:
            config (configparser.SectionProxy): Vessel section defining a 
              Vessel

        Raises:
            ValueError: Raised if section does not contain Address parameter

        Returns:
            classes.vessel.Vessel: Vessel object for the vessel specified in
              the config section
        """

        host = None
        username = None
        password = None
        database = None
        ssh = False
        ssh_username = None
        ssh_password = None
        ssh_timeout = 100
        ssh_passphrase = None

        if "Username" in config.keys():
            username = config["Username"]

        if "Password" in config.keys():
            password = config["Password"]

        if "Database" in config.keys():
            database = config["Database"]

        if "SSH" in config.keys():
            if int(config["SSH"]) == 1:
                ssh = True

        return cls(config.name.split()[1], config["Host"], username, password, database, ssh, ssh_username, ssh_password, ssh_timeout, ssh_passphrase)

    def __init__(self, name: str, host: str, username: Optional[str] = None,
            password: Optional[str] = None, database: Optional[str] = None,
            ssh = False, ssh_username = None, ssh_password = None, ssh_timeout = None, ssh_passphrase = None) -> None:
        """Initialize new Vessel object

        Args:
            name (str): Name of the Vessel
        """
        self.name = name
        self.host = host
        self.username = username
        self.password = password
        self.database = database
        self.ssh = ssh
        self.ssh_username = ssh_username
        self.ssh_password = ssh_password
        self.ssh_timeout = ssh_timeout
        self.ssh_passphrase = ssh_passphrase

        self.db = self.connect()

    def connect(self):
        return Database(self)

    def getCourses(self) -> list:
        results = self.db._execute(QUERY_COURSE, ctype=DictCursor)
        return results

    def getUserInfoFields(self) -> list:
        results = self.db._execute(QUERY_USER_INFO_FIELD, ctype=DictCursor)
        return results

    def getUserInfoData(self, field: Optional[int] = None, user: Optional[int] = None) -> list:
        query = QUERY_USER_INFO_DATA
        parameters = []

        if field:
            query += " WHERE fieldid = %s"
            parameters.append(int(field))

        if user:
            if query != QUERY_USER_INFO_DATA:
                query += " AND "
            else:
                query += " WHERE "

            query += "userid = %s"
            parameters.append(int(user))

        results = self.db._execute(query, tuple(parameters), ctype=DictCursor)

        return results

    def getUsers(self, username: Optional[str] = None, id: Optional[int] = None) -> dict:
        query = QUERY_USER
        parameters = tuple()

        if username:
            query += f" WHERE username = %s"
            parameters = (username,)

        elif id:
            query += f" WHERE id = %s"
            parameters = (int(id),)

        results = self.db._execute(query, parameters, ctype=DictCursor)

        users = dict()

        for result in results:
            user = result
            result["custom_fields"] = dict()
            users[user["id"]] = user

        ofields = self.getUserInfoFields()

        for ofield in ofields:
            odata = self.getUserInfoData(ofield["id"], id)

            for value in odata:
                users[value["userid"]]["custom_fields"][ofield["shortname"]] = value["data"]

        return users

    def getHTMLCerts(self, after: int = 0, before: int = int(datetime.now().timestamp())):
        results = self.db._execute(f"{QUERY_HTML_CERT_ISSUES} {QUERY_WHERE_TIMESTAMPS % {'column': 'timecreated', 'after': after, 'before': before}}", ctype=DictCursor)
        ocerts = self.db._execute(QUERY_HTML_CERT, ctype=DictCursor)

        certs = dict()

        for ocert in ocerts:
            certs[ocert["id"]] = ocert

        for result in results:
            try:
                result["cert"] = certs[result["htmlcertid"]]
            except KeyError:
                result["cert"] = None

            result["certtype"] = "htmlcert"

        return results

    def getCustomCerts(self, after: int = 0, before: int = int(datetime.now().timestamp())):
        results = self.db._execute(f"{QUERY_CUSTOM_CERT_ISSUES} {QUERY_WHERE_TIMESTAMPS % {'column': 'timecreated', 'after': after, 'before': before}}", ctype=DictCursor)
        ocerts = self.db._execute(QUERY_CUSTOM_CERT, ctype=DictCursor)

        certs = dict()

        for ocert in ocerts:
            certs[ocert["id"]] = ocert

        for result in results:
            try:
                result["cert"] = certs[result["customcertid"]]
            except KeyError:
                result["cert"] = None

            result["certtype"] = "customcert"

        return results

    def getCerts(self, after: int = 0, before: int = int(datetime.now().timestamp())):
        return sorted(self.getHTMLCerts(after, before) + self.getCustomCerts(after, before), key=lambda d: d["timecreated"])