From 081b9f74e916b05c74fb2ce9334a261bad9fe5bb Mon Sep 17 00:00:00 2001
From: CommanderKeynes <andrewjackson947@gmail.coma>
Date: Thu, 5 Sep 2024 20:59:26 -0500
Subject: [PATCH] Add tests

---
 tests/docker/Dockerfile  |  4 +-
 tests/python/test_krb.py | 94 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 97 insertions(+), 1 deletion(-)
 create mode 100644 tests/python/test_krb.py

diff --git a/tests/docker/Dockerfile b/tests/docker/Dockerfile
index 6eccddd..d50bb1d 100644
--- a/tests/docker/Dockerfile
+++ b/tests/docker/Dockerfile
@@ -2,7 +2,9 @@ FROM rust:bullseye
 
 COPY --from=sclevine/yj /bin/yj /bin/yj
 RUN /bin/yj -h
-RUN apt-get update && apt-get install llvm-11 psmisc postgresql-contrib postgresql-client ruby ruby-dev libpq-dev python3 python3-pip lcov curl sudo iproute2 -y
+RUN apt-get update && apt-get install llvm-11 psmisc postgresql-contrib postgresql-client ruby ruby-dev libpq-dev python3 python3-pip lcov curl sudo iproute2 gnupg lsb-release -y
+RUN env DEBIAN_FRONTEND=noninteractive apt-get -y install krb5-kdc krb5-admin-server krb5-user
+
 RUN cargo install cargo-binutils rustfilt
 RUN rustup component add llvm-tools-preview
 RUN sudo gem install bundler
diff --git a/tests/python/test_krb.py b/tests/python/test_krb.py
new file mode 100644
index 0000000..db460fa
--- /dev/null
+++ b/tests/python/test_krb.py
@@ -0,0 +1,94 @@
+import signal
+import socket
+import subprocess
+import utils
+
+REALM = 'EXAMPLE.COM'
+SUPPORTED_ENCRYPTION_TYPES = 'aes256-cts-hmac-sha1-96:normal'
+KADMIN_PRINCIPAL = 'root'
+KADMIN_PASSWORD = 'root'
+KDC_KADMIN_SERVER = socket.gethostname()
+
+LOGDIR = 'log'
+PG_LOG = f'{LOGDIR}/krb.log'
+# Assumes packages are installed; krb5-kdc and krb5-admin-server on debian
+KADMIN_PRINCIPAL_FULL = f'{KADMIN_PRINCIPAL}@{REALM}'
+MASTER_PASSWORD = 'master_password'
+
+
+def setup_krb():
+    krb5_conf = f"""
+[libdefaults]
+        default_realm = {REALM}
+        rdns = false
+
+[realms]
+        {REALM} = {{
+                kdc_ports = 88,750
+                kadmind_port = 749
+                kdc = {KDC_KADMIN_SERVER}
+                admin_server = {KDC_KADMIN_SERVER}
+        }}
+    """
+    with open("/etc/krb5.conf", "w") as text_file:
+        text_file.write(krb5_conf)
+
+    kdc_conf = f"""
+[realms]
+        {REALM} = {{
+                acl_file = /etc/krb5kdc/kadm5.acl
+                max_renewable_life = 7d 0h 0m 0s
+                supported_enctypes = {SUPPORTED_ENCRYPTION_TYPES}
+                default_principal_flags = +preauth
+        }}
+    """
+    with open("/etc/krb5kdc/kdc.conf", "w") as text_file:
+        text_file.write(kdc_conf)
+
+    kadm5_acl = f"""
+    {KADMIN_PRINCIPAL_FULL} *
+    """
+    with open("/etc/krb5kdc/kadm5.acl", "w") as text_file:
+        text_file.write(kadm5_acl)
+
+    kerberos_command = f"""
+    krb5_newrealm <<EOF
+    {MASTER_PASSWORD}
+    {MASTER_PASSWORD}
+    EOF
+    """
+    subprocess.run(kerberos_command, check=False, shell=True)
+
+    delete_principal = f'kadmin.local -q "delete_principal -force {KADMIN_PRINCIPAL_FULL}"'
+    subprocess.run(delete_principal, check=True, shell=True)
+
+    create_principal = f'kadmin.local -q "addprinc -pw {KADMIN_PASSWORD} {KADMIN_PRINCIPAL_FULL}"'
+    subprocess.run(create_principal, check=True, shell=True)
+
+    kinit_command = f'echo {KADMIN_PASSWORD} | kinit'
+    subprocess.run(kinit_command, check=True, shell=True)
+
+    utils.pgcat_start()
+
+
+def teardown_krb():
+    subprocess.run('kdestroy', check=True, shell=True)
+
+    delete_principal = f'kadmin.local -q "delete_principal -force {KADMIN_PRINCIPAL_FULL}"'
+    subprocess.run(delete_principal, check=True, shell=True)
+
+    utils.pg_cat_send_signal(signal.SIGINT)
+
+
+def test_krb():
+    setup_krb()
+    # TODO test connect to database
+
+    utils.pgcat_start()
+    conn, cur = utils.connect_db(autocommit=False)
+    cur.execute("SELECT 1")
+    res = cur.fetchall()
+    print(res)
+    utils.cleanup_conn(conn, cur)
+
+    teardown_krb()