simple_server.py

"""Basic server for password auth."""

import logging
import socket

import paramiko

# Set up logging globally.
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger()

sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock.bind(("", 2222))

HOSTKEY_ED25519 = paramiko.Ed25519Key(filename="keys/server_ed25519")


class Server(paramiko.ServerInterface):
    def check_channel_request(self, kind, chanid):
        if kind == "session":
            return paramiko.OPEN_SUCCEEDED

    def check_auth_password(self, username: str, password: str):
        if username == "user" and password == "password":
            return paramiko.AUTH_SUCCESSFUL
        elif username == "2fa_user" and password == "password":
            return paramiko.AUTH_SUCCESSFUL
        else:
            return paramiko.AUTH_FAILED

    def check_channel_exec_request(self, channel: paramiko.Channel, command: bytes):
        self.event.set()
        cmd = command.decode()
        logger.info("channel_exec request received: %s", cmd)

        if cmd == "banner":
            self.banner(channel)
        elif cmd.startswith("echo "):
            self.echo(channel, cmd)
        else:
            return False
        return True

    def banner(self, chan: paramiko.Channel):
        chan.send(
            """===== Test server banner =====
    We should include some ascii art...
    """
        )

    def echo(self, chan: paramiko.Channel, cmd: str):
        _, s = cmd.split("echo ")
        chan.send(s)

    def listen(self):
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind(("", 2222))
        sock.listen(100)
        client, addr = sock.accept()

        t = paramiko.Transport(client)
        t.set_gss_host(socket.getfqdn(""))
        t.load_server_moduli()

        t.add_server_key(HOSTKEY_ED25519)

        # Starts the server and negotiates a new session as server. Either returns
        # successfully or raises SSHException.
        t.start_server(server=self)

        chan = t.accept(20)
        if chan is None:
            raise ValueError("No channel")
        # Wait for client to send a command.
        self.event.wait(30.0)
        if not self.event.is_set():
            raise ValueError("Client did not send a command")

        chan.close()
        t.close()


if __name__ == "__main__":
    while True:
        try:
            server = Server()
            server.listen()
        except KeyboardInterrupt:
            exit(0)
        except OSError:
            logger.exception("Caught OSError, usually address already in use")
            exit(1)
        except Exception as exc:
            logger.error(exc)