from overseer.models import Base, ScanTarget, ScanHistory
import ipaddress
from datetime import datetime
from os import path
from sqlalchemy import create_engine
from sqlalchemy.orm import Session
from sqlalchemy.orm.exc import NoResultFound


class DatabaseConnector:
    """Used to manipulate the database

    Methods
    -------
    create_scan_target(**kwargs)
        Create a new ScanTarget database row with the provided hostname or
        ip_addr
    get_scan_target(**kwargs)
        Get a ScanTarget by either its hostname or ip_addr
    get_all_scan_targets(page=1)
        Get all ScanTargets ordered by most recent, limited to 25 / page
    create_scan_result(status, results=[], error=None, **kwargs)
        Create a new ScanHistory for the provided hostname or ip_addr
    update_scan_result(history_id, status, results=None, error=None)
        Update the referenced ScanHistory ID with the provided values
    get_scan_results_by_target(page=1, **kwargs)
        Get all ScanHistory for the provided hostname or ip_add, limited to
        25 / page
    """

    def __init__(self, db_type, data_path=None):
        """
        Parameters
        ----------
        db_type : str
            Database tyle (e.g. sqlite, memory)
        data_path : str
            Directory to store the sqlite file
        """
        if db_type.lower() == "memory":
            self.__engine = create_engine(
                "sqlite+pysqlite:///:memory:", echo=False, future=True
            )
        elif db_type.lower() == "sqlite":
            db_path = path.join(data_path, "overseer.sqlite")
            self.__engine = create_engine(
                "sqlite+pysqlite:///%s" % db_path,
                echo=False,
                future=True,
            )
        Base.metadata.create_all(self.__engine)
        self.__cleanup_stale_records()

    def __cleanup_stale_records(self):
        """Cleans up any stale ScanHistory records"""
        session = Session(bind=self.__engine)
        history_filter = ScanHistory.status == "IN_PROGRESS"
        all_stale = session.query(ScanHistory).filter(history_filter).all()

        for stale in all_stale:
            stale.status = "FAILED"
            stale.error = "Stale history"

        session.commit()
        session.close()

    def create_scan_target(self, **kwargs):
        """Create a new ScanTarget database row with the provided hostname or
        ip_addr

        Parameters
        ----------
        **kwargs
            Either hostname or ip_addr

        Raises
        ------
        ValueError
            If hostname or ip_addr isn't provided in kwargs

        Returns
        -------
        ScanTarget
            The created ScanTarget
        """
        if len(kwargs.keys() & {"ip_addr", "hostname"}) != 1:
            raise ValueError("Missing keyword argument: ip_addr or hostname")

        if "ip_addr" in kwargs:
            int_ip_addr = int(ipaddress.ip_address(kwargs["ip_addr"]))
            scan_target = ScanTarget(ip=int_ip_addr, updated_at=datetime.now())
        elif "hostname" in kwargs:
            scan_target = ScanTarget(
                hostname=kwargs["hostname"], updated_at=datetime.now()
            )

        session = Session(bind=self.__engine, expire_on_commit=False)
        session.add(scan_target)
        session.commit()
        session.close()
        return scan_target

    def get_scan_target(self, **kwargs):
        """Get a ScanTarget by either its hostname or ip_addr

        Parameters
        ----------
        **kwargs
            Either hostname or ip_addr

        Raises
        ------
        ValueError
            If hostname or ip_addr isn't provided in kwargs

        Returns
        -------
        ScanTarget
            The requested ScanTarget
        """
        if len(kwargs.keys() & {"ip_addr", "hostname"}) != 1:
            raise ValueError("Missing keyword argument: ip_addr or hostname")

        if "ip_addr" in kwargs:
            int_ip_addr = int(ipaddress.ip_address(kwargs["ip_addr"]))
            target_filter = ScanTarget.ip == int_ip_addr
        elif "hostname" in kwargs:
            target_filter = ScanTarget.hostname == kwargs["hostname"]

        session = Session(bind=self.__engine)
        scan_target = session.query(ScanTarget).filter(target_filter).first()
        session.close()
        return scan_target

    def get_all_scan_targets(self, page=1):
        """Get all ScanTargets ordered by most recent, limited to 25 / page

        Parameters
        ----------
        page : int, optional
            The desired ScanTarget page

        Returns
        -------
        list
            List of ScanTarget
        """
        session = Session(bind=self.__engine)
        scan_targets = (
            session.query(ScanTarget)
            .order_by(ScanTarget.updated_at.desc())
            .offset((page - 1) * 25)
            .limit(25)
            .all()
        )
        session.close()
        return scan_targets

    def create_scan_result(self, status, results=[], error=None, **kwargs):
        """Create a new ScanHistory for the provided hostname or ip_addr

        Parameters
        ----------
        status : str
            The status of the scan (IN_PROGRESS, FAILED, COMPLETE)
        results : list
            List of strings of open ports (E.g. ["53 UDP", "53 TCP"]
        error : str, optional
            Error message, if any
        **kwargs
            Either hostname or ip_addr

        Returns
        -------
        ScanHistory
            The created ScanHistory
        """
        scan_target = self.get_scan_target(**kwargs)
        if not scan_target:
            scan_target = self.create_scan_target(**kwargs)

        scan_history = ScanHistory(
            target_id=scan_target.id,
            results=",".join(results),
            status=status,
            error=error,
            created_at=datetime.now(),
        )
        session = Session(bind=self.__engine, expire_on_commit=False)
        session.add(scan_history)
        session.commit()
        session.close()
        return scan_history

    def update_scan_result(self, history_id, status, results=None, error=None):
        """Update the referenced ScanHistory ID with the provided values

        Parameters
        ----------
        history_id : int
            The ScanHistory ID to update
        status : str
            The status of the scan (IN_PROGRESS, FAILED, COMPLETE)
        results : list, optional
            List of strings of open ports (E.g. ["53 UDP", "53 TCP"]
        error : str, optional
            Error message, if any

        Raises
        ------
        NoResultFound
            If we cannot find the desired ScanHistory by ID

        Returns
        -------
        ScanHistory
            The updated ScanHistory
        """
        session = Session(bind=self.__engine, expire_on_commit=False)
        scan_history = session.query(ScanHistory).get(history_id)

        if scan_history is None:
            raise NoResultFound("ScanHistory %s does not exist" % history_id)

        scan_history.status = status
        if results is not None:
            scan_history.results = ",".join(results)
        if error is not None:
            scan_history.error = error

        session.commit()
        session.close()
        return scan_history

    def get_scan_results_by_target(self, page=1, **kwargs):
        """Get all ScanHistory for the provided hostname or ip_add, limited to
        25 / page

        Parameters
        ----------
        page : int, optional
            The desired ScanResult page

        Returns
        -------
        list
            List of ScanHistory
        """
        if len(kwargs.keys() & {"ip_addr", "hostname"}) != 1:
            raise ValueError("Missing keyword argument: ip_addr or hostname")

        if "ip_addr" in kwargs:
            int_ip_addr = int(ipaddress.ip_address(kwargs["ip_addr"]))
            history_filter = ScanTarget.ip == int_ip_addr
        elif "hostname" in kwargs:
            history_filter = ScanTarget.hostname == kwargs["hostname"]

        session = Session(bind=self.__engine)
        scan_histories = (
            session.query(ScanHistory)
            .join(ScanHistory.target)
            .filter(history_filter)
            .order_by(ScanHistory.created_at.desc())
            .offset((page - 1) * 25)
            .limit(25)
            .all()
        )
        session.close()
        return scan_histories
