From 16e54218bdf148f67f0449f0846815b736c4918f Mon Sep 17 00:00:00 2001
From: Dmitriy Pavlov <dpavlov@smartling.com>
Date: Wed, 12 Apr 2023 17:28:05 -0400
Subject: [PATCH] TLS check: implement StartTLS protocol for Mysql (#14352)

* TLS check: implement StartTLS protocol for Mysql

* fix model validation
---
 tls/assets/configuration/spec.yaml            |  3 ++-
 .../tls/config_models/instance.py             |  2 +-
 tls/datadog_checks/tls/data/conf.yaml.example |  2 +-
 tls/datadog_checks/tls/tls_remote.py          | 23 ++++++++++++++++++-
 tls/tests/compose/Dockerfile                  |  5 ++++
 tls/tests/compose/docker-compose.yml          | 14 +++++++++++
 tls/tests/conftest.py                         | 13 ++++++++++-
 tls/tests/test_remote.py                      | 14 +++++++++++
 8 files changed, 71 insertions(+), 5 deletions(-)

diff --git a/tls/assets/configuration/spec.yaml b/tls/assets/configuration/spec.yaml
index 67091e9641282..7a69f99ad14a5 100644
--- a/tls/assets/configuration/spec.yaml
+++ b/tls/assets/configuration/spec.yaml
@@ -140,10 +140,11 @@ files:
         irc, postgres, mysql, lmtp, nntp, sieve, and ldap.
 
         Currently this checks supports only the below protocols:
-          postgres
+          postgres, mysql
       value:
         type: string
         enum:
         - postgres
+        - mysql
     - template: instances/default
     - template: instances/tls
diff --git a/tls/datadog_checks/tls/config_models/instance.py b/tls/datadog_checks/tls/config_models/instance.py
index 661250a77bbdb..9a37b81081aec 100644
--- a/tls/datadog_checks/tls/config_models/instance.py
+++ b/tls/datadog_checks/tls/config_models/instance.py
@@ -49,7 +49,7 @@ class Config:
     server: str
     server_hostname: Optional[str]
     service: Optional[str]
-    start_tls: Optional[Literal['postgres']]
+    start_tls: Optional[Literal['postgres', 'mysql']]
     tags: Optional[Sequence[str]]
     timeout: Optional[int]
     tls_ca_cert: Optional[str]
diff --git a/tls/datadog_checks/tls/data/conf.yaml.example b/tls/datadog_checks/tls/data/conf.yaml.example
index 52618761dd14b..da1f08328f4bd 100644
--- a/tls/datadog_checks/tls/data/conf.yaml.example
+++ b/tls/datadog_checks/tls/data/conf.yaml.example
@@ -133,7 +133,7 @@ instances:
     ## irc, postgres, mysql, lmtp, nntp, sieve, and ldap.
     ##
     ## Currently this checks supports only the below protocols:
-    ##   postgres
+    ##   postgres, mysql
     #
     # start_tls: <START_TLS>
 
diff --git a/tls/datadog_checks/tls/tls_remote.py b/tls/datadog_checks/tls/tls_remote.py
index 9a61c8ad76a04..96d2cbf888cde 100644
--- a/tls/datadog_checks/tls/tls_remote.py
+++ b/tls/datadog_checks/tls/tls_remote.py
@@ -3,7 +3,7 @@
 # Licensed under a 3-clause BSD style license (see LICENSE)
 import ssl
 from hashlib import sha256
-from struct import pack
+from struct import pack, unpack
 
 from cryptography.x509.base import load_der_x509_certificate
 from cryptography.x509.extensions import ExtensionNotFound
@@ -134,6 +134,27 @@ def _switch_starttls(self, sock):
             data = self._read_n_bytes_from_socket(sock, 1)
             if data != b'S':
                 raise Exception('Postgres endpoint does not support TLS')
+        elif protocol == "mysql":
+            self.log.debug('Switching connection to encrypted for %s protocol', protocol)
+            cap_protocol_41 = 1 << 9
+            cap_ssl = 1 << 11
+            cap_secure_connection = 1 << 15
+            capabilities = cap_protocol_41 | cap_ssl | cap_secure_connection
+            max_packet_len = 2**24 - 1
+            charset_id = 8  # latin1
+            # Form Protocol::SSLRequest packet
+            data_init = pack("<iIB23s", capabilities, max_packet_len, charset_id, b"")
+            # Form Mysql Protocol::Packet
+            packet_len = pack("<I", len(data_init))[:3]
+            packet_seq = pack("<B", 1)
+            packet = packet_len + packet_seq + data_init
+            # Read 4 bytes of header to get packet length
+            packet_header = self._read_n_bytes_from_socket(sock, 4)
+            btrl, btrh, packet_number = unpack("<HBB", packet_header)
+            bytes_to_read = btrl + (btrh << 16)
+            # Read Mysql welcome message
+            data = self._read_n_bytes_from_socket(sock, bytes_to_read)
+            sock.sendall(packet)
         else:
             raise Exception('Unsupported starttls protocol: ' + protocol)
 
diff --git a/tls/tests/compose/Dockerfile b/tls/tests/compose/Dockerfile
index 2c524a857dcff..db839658a4bfa 100644
--- a/tls/tests/compose/Dockerfile
+++ b/tls/tests/compose/Dockerfile
@@ -36,3 +36,8 @@ RUN openssl genrsa -out expired.mock.key 2048 \
 ENV PSQL_UID=70
 ENV PSQL_GUID=70
 RUN for file in valid.mock.key valid.mock.crt; do cp $file postgres.$file; chown $PSQL_UID:$PSQL_GUID postgres.$file; done
+
+## Prepare mysql certificates
+ENV MYSQL_UID=999
+ENV MYSQL_GUID=999
+RUN for file in valid.mock.key valid.mock.crt; do cp $file mysql.$file; chown $MYSQL_UID:$MYSQL_GUID mysql.$file; done
diff --git a/tls/tests/compose/docker-compose.yml b/tls/tests/compose/docker-compose.yml
index b3c14ae5d425b..9796151bd08c2 100644
--- a/tls/tests/compose/docker-compose.yml
+++ b/tls/tests/compose/docker-compose.yml
@@ -58,5 +58,19 @@ services:
     depends_on:
       - cert-builder
 
+  # mysql serving valid.mysql.mock
+  starttls-mysql:
+    container_name: mysql-valid
+    image: mysql
+    volumes:
+      - certs:/var/lib/certs/:ro
+    ports:
+      - "3306:3306"
+    command: --default-authentication-plugin=mysql_native_password --ssl-cert=/var/lib/certs/mysql.valid.mock.crt --ssl-key=/var/lib/certs/mysql.valid.mock.key
+    environment:
+      - MYSQL_ROOT_PASSWORD=MysqlR00tpass
+    depends_on:
+      - cert-builder
+
 volumes:
   certs:
diff --git a/tls/tests/conftest.py b/tls/tests/conftest.py
index b608121f3dc57..9247cf8a19923 100644
--- a/tls/tests/conftest.py
+++ b/tls/tests/conftest.py
@@ -32,7 +32,7 @@
 
 @pytest.fixture(scope='session', autouse=True)
 def dd_environment(instance_e2e, mock_local_tls_dns):
-    with docker_run(os.path.join(HERE, 'compose', 'docker-compose.yml'), build=True, sleep=5):
+    with docker_run(os.path.join(HERE, 'compose', 'docker-compose.yml'), build=True, sleep=20):
         e2e_metadata = {'docker_volumes': ['{}:{}'.format(CA_CERT, CA_CERT_MOUNT_PATH)]}
         yield instance_e2e, e2e_metadata
 
@@ -258,3 +258,14 @@ def instance_remote_postgresql_valid():
         'start_tls': 'postgres',
         'tls_ca_cert': CA_CERT,
     }
+
+
+@pytest.fixture
+def instance_remote_mysql_valid():
+    return {
+        'server': 'localhost',
+        'port': 3306,
+        'server_hostname': 'valid.mock',
+        'start_tls': 'mysql',
+        'tls_ca_cert': CA_CERT,
+    }
diff --git a/tls/tests/test_remote.py b/tls/tests/test_remote.py
index bea8a2e546fe1..3c57cb56154c7 100644
--- a/tls/tests/test_remote.py
+++ b/tls/tests/test_remote.py
@@ -321,3 +321,17 @@ def test_postgres_ok(aggregator, instance_remote_postgresql_valid):
     aggregator.assert_metric('tls.days_left', count=1)
     aggregator.assert_metric('tls.seconds_left', count=1)
     aggregator.assert_all_metrics_covered()
+
+
+def test_mysql_ok(aggregator, instance_remote_mysql_valid):
+    c = TLSCheck('tls', {}, [instance_remote_mysql_valid])
+    c.check(None)
+
+    aggregator.assert_service_check(SERVICE_CHECK_CAN_CONNECT, status=c.OK, tags=c._tags, count=1)
+    aggregator.assert_service_check(SERVICE_CHECK_VERSION, status=c.OK, tags=c._tags, count=1)
+    aggregator.assert_service_check(SERVICE_CHECK_VALIDATION, status=c.OK, tags=c._tags, count=1)
+    aggregator.assert_service_check(SERVICE_CHECK_EXPIRATION, status=c.OK, tags=c._tags, count=1)
+
+    aggregator.assert_metric('tls.days_left', count=1)
+    aggregator.assert_metric('tls.seconds_left', count=1)
+    aggregator.assert_all_metrics_covered()