import base64
import socket
import ssl
import sys

def wait_for(sock, expect):
    msg = b""
    i = 0
    while expect not in msg:
        msg += sock.recv(1024)
        i += 1
        if i == 50:
            raise Exception(f"{expect} not in {msg}")

    print("<", msg.decode("utf-8").strip("\r\n"))

def write(sock, msg):
    print(">", msg.decode("utf-8").strip("\r\n"))
    sock.sendall(msg)


hostname = sys.argv[1]
password = sys.argv[2]

context = ssl.create_default_context()

with socket.create_connection((hostname, 6697)) as sock:
    with context.wrap_socket(sock, server_hostname=hostname) as ssock:
        write(ssock, b"CAP REQ :sasl\r\n")
        wait_for(ssock, b"CAP * ACK sasl\r\n")
        write(ssock, b"NICK autopkgtest\r\n")
        write(ssock, b"USER autopkgtest 0 * :Autopkgtest user\r\n")
        write(ssock, b"AUTHENTICATE plain\r\n")
        wait_for(ssock, b"AUTHENTICATE +\r\n")
        auth_data = base64.b64encode(f"\0autopkgtest\0{password}".encode("utf-8"))  # per rfc4616
        write(ssock, b"AUTHENTICATE " + auth_data + b"\r\n")
        wait_for(ssock, b"SASL authentication successful\r\n")
        write(ssock, b"CAP END\r\n")
        wait_for(ssock, b"001 autopkgtest :Welcome to soju, autopkgtest\r\n")
