# Copyright (c) 2009-2014 Upi Tamminen <desaster@gmail.com>
# See the COPYRIGHT file for more information

"""
This module contains ...
"""

from __future__ import annotations

from configparser import NoOptionError
import time
from typing import Optional

from twisted.conch.openssh_compat import primes
from twisted.conch.ssh import factory, keys, transport
from twisted.cred import portal as tp
from twisted.python import log

from cowrie.core.config import CowrieConfig
from cowrie.ssh import connection
from cowrie.ssh import keys as cowriekeys
from cowrie.ssh import transport as shellTransport
from cowrie.ssh.userauth import HoneyPotSSHUserAuthServer
from cowrie.ssh_proxy import server_transport as proxyTransport
from cowrie.ssh_proxy.userauth import ProxySSHAuthServer


class CowrieSSHFactory(factory.SSHFactory):
    """
    This factory creates HoneyPotSSHTransport instances
    They listen directly to the TCP port
    """

    starttime: Optional[float] = None
    privateKeys: dict[bytes, bytes] = {}
    publicKeys: dict[bytes, bytes] = {}
    primes = None
    portal: Optional[tp.Portal] = None  # gets set by plugin
    ourVersionString: bytes = CowrieConfig.get(
        "ssh", "version", fallback="SSH-2.0-OpenSSH_6.0p1 Debian-4+deb7u2"
    ).encode("ascii")

    def __init__(self, backend, pool_handler):
        self.pool_handler = pool_handler
        self.backend: str = backend
        self.services = {
            b"ssh-userauth": ProxySSHAuthServer
            if self.backend == "proxy"
            else HoneyPotSSHUserAuthServer,
            b"ssh-connection": connection.CowrieSSHConnection,
        }
        super().__init__()

    def logDispatch(self, **args):
        """
        Special delivery to the loggers to avoid scope problems
        """
        args["sessionno"] = "S{}".format(args["sessionno"])
        for output in self.tac.output_plugins:
            output.logDispatch(**args)

    def startFactory(self):
        # For use by the uptime command
        self.starttime = time.time()

        # Load/create keys
        try:
            public_key_auth = [
                i.encode("utf-8")
                for i in CowrieConfig.get("ssh", "public_key_auth").split(",")
            ]
        except NoOptionError:
            # no keys defined, use the three most common pub keys of OpenSSH
            public_key_auth = [b"ssh-rsa", b"ecdsa-sha2-nistp256", b"ssh-ed25519"]
        for key in public_key_auth:
            if key == b"ssh-rsa":
                rsaPubKeyString, rsaPrivKeyString = cowriekeys.getRSAKeys()
                self.publicKeys[key] = keys.Key.fromString(data=rsaPubKeyString)
                self.privateKeys[key] = keys.Key.fromString(data=rsaPrivKeyString)
            elif key == b"ssh-dss":
                dsaaPubKeyString, dsaPrivKeyString = cowriekeys.getDSAKeys()
                self.publicKeys[key] = keys.Key.fromString(data=dsaaPubKeyString)
                self.privateKeys[key] = keys.Key.fromString(data=dsaPrivKeyString)
            elif key == b"ecdsa-sha2-nistp256":
                ecdsaPuKeyString, ecdsaPrivKeyString = cowriekeys.getECDSAKeys()
                self.publicKeys[key] = keys.Key.fromString(data=ecdsaPuKeyString)
                self.privateKeys[key] = keys.Key.fromString(data=ecdsaPrivKeyString)
            elif key == b"ssh-ed25519":
                ed25519PubKeyString, ed25519PrivKeyString = cowriekeys.geted25519Keys()
                self.publicKeys[key] = keys.Key.fromString(data=ed25519PubKeyString)
                self.privateKeys[key] = keys.Key.fromString(data=ed25519PrivKeyString)

        _modulis = "/etc/ssh/moduli", "/private/etc/moduli"
        for _moduli in _modulis:
            try:
                self.primes = primes.parseModuliFile(_moduli)
                break
            except OSError:
                pass

        # this can come from backend in the future, check HonSSH's slim client
        self.ourVersionString = CowrieConfig.get(
            "ssh", "version", fallback="SSH-2.0-OpenSSH_6.0p1 Debian-4+deb7u2"
        ).encode("ascii")

        factory.SSHFactory.startFactory(self)
        log.msg("Ready to accept SSH connections")

    def stopFactory(self):
        factory.SSHFactory.stopFactory(self)

    def buildProtocol(self, addr):
        """
        Create an instance of the server side of the SSH protocol.

        @type addr: L{twisted.internet.interfaces.IAddress} provider
        @param addr: The address at which the server will listen.

        @rtype: L{cowrie.ssh.transport.HoneyPotSSHTransport}
        @return: The built transport.
        """
        t: transport.SSHServerTransport
        if self.backend == "proxy":
            t = proxyTransport.FrontendSSHTransport()
        else:
            t = shellTransport.HoneyPotSSHTransport()

        t.ourVersionString = self.ourVersionString
        t.supportedPublicKeys = list(self.privateKeys.keys())

        if not self.primes:
            ske = t.supportedKeyExchanges[:]
            if b"diffie-hellman-group-exchange-sha1" in ske:
                ske.remove(b"diffie-hellman-group-exchange-sha1")
                log.msg("No moduli, no diffie-hellman-group-exchange-sha1")
            if b"diffie-hellman-group-exchange-sha256" in ske:
                ske.remove(b"diffie-hellman-group-exchange-sha256")
                log.msg("No moduli, no diffie-hellman-group-exchange-sha256")
            t.supportedKeyExchanges = ske

        try:
            t.supportedCiphers = [
                i.encode("utf-8") for i in CowrieConfig.get("ssh", "ciphers").split(",")
            ]
        except NoOptionError:
            # Reorder supported ciphers to resemble current openssh more
            t.supportedCiphers = [
                b"aes128-ctr",
                b"aes192-ctr",
                b"aes256-ctr",
                b"aes256-cbc",
                b"aes192-cbc",
                b"aes128-cbc",
                b"3des-cbc",
                b"blowfish-cbc",
                b"cast128-cbc",
            ]

        try:
            t.supportedMACs = [
                i.encode("utf-8") for i in CowrieConfig.get("ssh", "macs").split(",")
            ]
        except NoOptionError:
            # SHA1 and MD5 are considered insecure now. Use better algos
            # like SHA-256 and SHA-384
            t.supportedMACs = [
                b"hmac-sha2-512",
                b"hmac-sha2-384",
                b"hmac-sha2-256",
                b"hmac-sha1",
                b"hmac-md5",
            ]

        try:
            t.supportedCompressions = [
                i.encode("utf-8")
                for i in CowrieConfig.get("ssh", "compression").split(",")
            ]
        except NoOptionError:
            t.supportedCompressions = [b"zlib@openssh.com", b"zlib", b"none"]

        t.factory = self

        return t
