import paramiko

"""
Provides a raw SSH session with an NGFW, useful for scripting checks against
a user session
"""


class SSHSession:

    BUFFER_LEN = 9999

    def __init__(self, device, username, password, echo=False):
        self.device = device
        self.username = username
        self.password = password
        self.echo = echo
        self._connect()
        self._get_prompt()
        self.mode = "op"
        self.error_text = ""

    def _connect(self):
        self.client = paramiko.client.SSHClient()
        self.client.load_system_host_keys()
        self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.client.connect(
            self.device,
            username=self.username,
            password=self.password
        )
        self.channel = self.client.invoke_shell()

    def _get_prompt(self):
        recvd = ""
        while ">" not in recvd:
            recvd = self.channel.recv(self.BUFFER_LEN).decode()
            if self.echo:
                print(recvd)
        self.prompt = recvd.strip()

    def get_error_text(self):
        """Get error text generated by this module"""
        return self.error_text

    def config_mode(self):
        """Enter configuration mode"""

        self.prompt = self.prompt.replace(">", "#")
        self.channel.send("configure\n")
        recvd = ""
        while self.prompt not in recvd:
            recvd = self.channel.recv(self.BUFFER_LEN).decode()
            if self.echo:
                print(recvd)
        self.mode = "config"

    def set_command(self, cmd):
        """Execute a set command inside config mode, returns False on error"""

        if not self.mode == "config":
            self.config_mode()
        self.channel.send(f"{cmd}\n")
        total_received = ""
        received = ""

        # Scan the output until the prompt is found
        while "[edit]" not in total_received:
            received = self.channel.recv(self.BUFFER_LEN).decode()
            if self.echo:
                print(received)
            total_received += received
            received = ""

        # Check the output for any errors
        error_text = ['Unknown command:', 'Invalid syntax.']
        for et in error_text:
            if et in total_received:
                self.error_text = total_received
                return False
        return True
