#!/usr/bin/env python3

"""Minimal fake MySQL server for the TLS verification test case.

Implements just enough of the MySQL client/server protocol to let a client
connect over TLS: the initial handshake, the SSLRequest packet, and an OK
packet for the authentication. It exists so that the test case can use a
certificate signed by a throwaway CA without setting up a real server.

Without --cert and --key it generates a throwaway certificate for itself, so
it can be run on its own. --cert-mode picks what it then presents, which is
the whole point of this test case:

  leaf        a certificate signed by a generated CA, sent WITHOUT that CA.
              This is the shape Connector/C 3.4 fails to verify, so it is the
              default.
  fullchain   the same certificate, sent together with the CA that issued it.
  selfsigned  a directly self-signed certificate, i.e. what a MariaDB 11.4+
              server presents when it has none configured.

The generated CA is written out as well, so it can be passed to a client with
--ssl-ca to check that the certificate is otherwise fine.

With --no-ssl the server doesn't announce CLIENT_SSL, i.e. it behaves like a
server that has no TLS support at all. No certificate is needed then.

Usage: fake-mysqld.py [--cert-mode <mode>] [--port <port>]
       fake-mysqld.py --cert <cert.pem> --key <key.pem> [--port <port>]
       fake-mysqld.py --no-ssl [--port <port>]
"""

import argparse
import atexit
import ipaddress
import os
import shutil
import socketserver
import ssl
import struct
import subprocess
import sys
import tempfile

CLIENT_LONG_PASSWORD = 0x00000001
CLIENT_LONG_FLAG = 0x00000004
CLIENT_CONNECT_WITH_DB = 0x00000008
CLIENT_PROTOCOL_41 = 0x00000200
CLIENT_SSL = 0x00000800
CLIENT_TRANSACTIONS = 0x00002000
CLIENT_SECURE_CONNECTION = 0x00008000
CLIENT_PLUGIN_AUTH = 0x00080000

SERVER_STATUS_AUTOCOMMIT = 0x0002
COM_QUIT = 0x01
COM_QUERY = 0x03

# Every query is answered with this, whatever was asked. A client that ends up
# here got a server it didn't authenticate.
QUERY_ANSWER = b"ATTACKER-CONTROLLED-DATA"

SERVER_VERSION = b"5.7.42-fake"
AUTH_PLUGIN = b"mysql_native_password"
# The scramble is never verified, any password is accepted.
SCRAMBLE = b"12345678" + b"123456789012"


class Connection:
    def __init__(self, sock, args):
        self.sock = sock
        self.args = args
        self.seq = 0
        self.tls = False

    def send_packet(self, payload):
        header = struct.pack("<I", len(payload))[:3] + struct.pack("<B", self.seq)
        self.sock.sendall(header + payload)
        self.seq = (self.seq + 1) % 256

    def read_exactly(self, count):
        data = b""
        while len(data) < count:
            more = self.sock.recv(count - len(data))
            if not more:
                return None
            data += more
        return data

    def read_packet(self):
        header = self.read_exactly(4)
        if header is None:
            return None
        length = struct.unpack("<I", header[:3] + b"\x00")[0]
        self.seq = (header[3] + 1) % 256
        return self.read_exactly(length)

    def send_handshake(self):
        caps = (CLIENT_LONG_PASSWORD | CLIENT_LONG_FLAG |
                CLIENT_CONNECT_WITH_DB | CLIENT_PROTOCOL_41 |
                CLIENT_TRANSACTIONS | CLIENT_SECURE_CONNECTION |
                CLIENT_PLUGIN_AUTH)
        if not self.args.no_ssl:
            caps |= CLIENT_SSL
        payload = b"\x0a" + SERVER_VERSION + b"\x00"
        payload += struct.pack("<I", 1)
        payload += SCRAMBLE[:8] + b"\x00"
        payload += struct.pack("<H", caps & 0xFFFF)
        payload += struct.pack("<B", 45)  # utf8mb4_general_ci
        payload += struct.pack("<H", SERVER_STATUS_AUTOCOMMIT)
        payload += struct.pack("<H", (caps >> 16) & 0xFFFF)
        payload += struct.pack("<B", len(SCRAMBLE) + 1)
        payload += b"\x00" * 10
        payload += SCRAMBLE[8:] + b"\x00"
        payload += AUTH_PLUGIN + b"\x00"
        self.send_packet(payload)

    def send_eof(self):
        self.send_packet(b"\xfe" + struct.pack("<HH", 0,
                                               SERVER_STATUS_AUTOCOMMIT))

    def send_query_answer(self):
        """Answer any query with a single row holding QUERY_ANSWER."""
        def lenenc_str(data):
            return struct.pack("<B", len(data)) + data  # len < 251 here

        self.send_packet(b"\x01")  # one column
        column = (lenenc_str(b"def") + lenenc_str(b"") + lenenc_str(b"") +
                  lenenc_str(b"") + lenenc_str(b"answer") +
                  lenenc_str(b"answer"))
        column += b"\x0c"
        column += struct.pack("<H", 45)     # utf8mb4_general_ci
        column += struct.pack("<I", 255)    # column length
        column += struct.pack("<B", 0xfd)   # MYSQL_TYPE_VAR_STRING
        column += struct.pack("<H", 0)      # flags
        column += struct.pack("<B", 0)      # decimals
        column += b"\x00\x00"               # filler
        self.send_packet(column)
        self.send_eof()
        self.send_packet(lenenc_str(QUERY_ANSWER))
        self.send_eof()

    def send_ok(self):
        payload = b"\x00\x00\x00"
        payload += struct.pack("<H", SERVER_STATUS_AUTOCOMMIT)
        payload += struct.pack("<H", 0)
        self.send_packet(payload)

    def handle(self):
        self.send_handshake()

        packet = self.read_packet()
        if packet is None:
            return
        client_caps = struct.unpack("<I", packet[:4])[0]
        if (client_caps & CLIENT_SSL) != 0 and not self.args.no_ssl:
            # The packet was an SSLRequest, the TLS handshake follows.
            context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
            context.load_cert_chain(self.args.cert, self.args.key)
            self.sock = context.wrap_socket(self.sock, server_side=True)
            if self.read_packet() is None:
                return
            self.tls = True
        print("connection: tls=%s" % ("yes" if self.tls else "no"), flush=True)
        self.send_ok()

        while True:
            packet = self.read_packet()
            if packet is None or packet[0] == COM_QUIT:
                return
            if packet[0] == COM_QUERY and \
                    packet[1:].lstrip().lower().startswith(b"select"):
                self.send_query_answer()
            else:
                # SET, USE and the like expect an OK packet, not a result set.
                self.send_ok()


class Handler(socketserver.BaseRequestHandler):
    def handle(self):
        try:
            Connection(self.request, self.server.args).handle()
        except (OSError, ssl.SSLError, struct.error):
            pass


class Server(socketserver.ThreadingTCPServer):
    allow_reuse_address = True
    daemon_threads = True

    def __init__(self, address, args):
        self.args = args
        super().__init__(address, Handler)


def run_openssl(*args):
    """Run openssl, and complain in a way that says what is missing."""
    try:
        subprocess.run(("openssl",) + args, check=True,
                       stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
    except FileNotFoundError:
        sys.exit("openssl is needed to generate a certificate, or pass "
                 "--cert and --key")
    except subprocess.CalledProcessError as e:
        sys.exit("openssl %s failed: %s" %
                 (args[0], e.stderr.decode(errors="replace").strip()))


def san_for(host):
    """subjectAltName for a host, as an IP address if it looks like one."""
    try:
        ipaddress.ip_address(host)
    except ValueError:
        return "DNS:" + host
    return "IP:" + host


def generate_certs(directory, host, mode):
    """Create a throwaway certificate for host in directory.

    Returns (cert file, key file, CA file). The CA file is None for a
    self-signed certificate, which has no CA.
    """
    key = os.path.join(directory, "server.key")
    cert = os.path.join(directory, "server.pem")
    san = san_for(host)

    if mode == "selfsigned":
        run_openssl("req", "-x509", "-newkey", "rsa:2048", "-nodes",
                    "-keyout", key, "-out", cert, "-days", "2",
                    "-subj", "/CN=" + host, "-addext", "subjectAltName=" + san)
        return cert, key, None

    ca_key = os.path.join(directory, "ca.key")
    ca_cert = os.path.join(directory, "ca.pem")
    run_openssl("req", "-x509", "-newkey", "rsa:2048", "-nodes",
                "-keyout", ca_key, "-out", ca_cert, "-days", "2",
                "-subj", "/CN=fake-mysqld throwaway CA")

    csr = os.path.join(directory, "server.csr")
    ext = os.path.join(directory, "server.ext")
    with open(ext, "w", encoding="utf-8") as f:
        f.write("subjectAltName=%s\n" % san)
    run_openssl("req", "-newkey", "rsa:2048", "-nodes", "-keyout", key,
                "-out", csr, "-subj", "/CN=" + host)
    leaf = os.path.join(directory, "leaf.pem")
    run_openssl("x509", "-req", "-in", csr, "-CA", ca_cert, "-CAkey", ca_key,
                "-CAcreateserial", "-out", leaf, "-days", "2",
                "-extfile", ext)

    with open(cert, "wb") as out:
        with open(leaf, "rb") as f:
            out.write(f.read())
        if mode == "fullchain":
            # Appending the CA is what makes OpenSSL report error 19 instead
            # of 20, and the certificate is then correctly refused.
            with open(ca_cert, "rb") as f:
                out.write(f.read())
    return cert, key, ca_cert


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=0)
    parser.add_argument("--cert")
    parser.add_argument("--key")
    parser.add_argument("--cert-mode", default="leaf",
                        choices=("leaf", "fullchain", "selfsigned"),
                        help="What to present when no --cert is given "
                             "(default: leaf)")
    parser.add_argument("--cert-host", default="localhost",
                        help="Name to issue the generated certificate for "
                             "(default: localhost)")
    parser.add_argument("--cert-dir",
                        help="Keep the generated certificates here instead "
                             "of in a temporary directory")
    parser.add_argument("--no-ssl", action="store_true",
                        help="Don't announce CLIENT_SSL, i.e. no TLS support")
    parser.add_argument("--port-file", help="Write the listening port here")
    args = parser.parse_args()
    if (args.cert is None) != (args.key is None):
        parser.error("--cert and --key go together")

    ca_cert = None
    generated = not args.no_ssl and args.cert is None
    if generated:
        if args.cert_dir is not None:
            directory = args.cert_dir
            os.makedirs(directory, exist_ok=True)
        else:
            directory = tempfile.mkdtemp(prefix="fake-mysqld-")
            atexit.register(shutil.rmtree, directory, True)
        args.cert, args.key, ca_cert = \
            generate_certs(directory, args.cert_host, args.cert_mode)

    server = Server((args.host, args.port), args)
    port = server.server_address[1]
    if args.port_file is not None:
        with open(args.port_file, "w", encoding="utf-8") as f:
            f.write("%u\n" % port)
    print("Listening on %s:%u (tls %s)" %
          (args.host, port, "disabled" if args.no_ssl else "enabled"),
          flush=True)
    if generated:
        if ca_cert is not None:
            print("Generated a %s certificate for %s, its CA is %s" %
                  (args.cert_mode, args.cert_host, ca_cert), flush=True)
        else:
            print("Generated a %s certificate for %s" %
                  (args.cert_mode, args.cert_host), flush=True)

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()