From 7898c49d51e7d090d0e94f89593c4dd461bfda41 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Romain=20Tarti=C3=A8re?= <romain@blogreen.org>
Date: Thu, 1 Dec 2022 15:21:38 -1000
Subject: [PATCH] Add support for MySQL and PostgreSQL handshakes

---
 lib/riemann/tools/tls_check.rb | 36 ++++++++++++++++++++++++++++++----
 1 file changed, 32 insertions(+), 4 deletions(-)

diff --git a/lib/riemann/tools/tls_check.rb b/lib/riemann/tools/tls_check.rb
index 8e51ba0b..bf1faa3b 100644
--- a/lib/riemann/tools/tls_check.rb
+++ b/lib/riemann/tools/tls_check.rb
@@ -8,10 +8,10 @@
 
 module URI
   {
-    'IMAP'       => 143,
-    'IMAPS'      => 993,
-    'MYSQL'      => 3306,
-    'POSTGRESQL' => 5432,
+    'IMAP'     => 143,
+    'IMAPS'    => 993,
+    'MYSQL'    => 3306,
+    'POSTGRES' => 5432,
   }.each do |scheme, port|
     klass = Class.new(Generic)
     klass.const_set('DEFAULT_PORT', port)
@@ -233,11 +233,39 @@ def tls_socket(uri, address)
           imap_tls_socket(uri, address)
         when 'ldap'
           ldap_tls_socket(uri, address)
+        when 'mysql'
+          mysql_tls_socket(uri, address)
+        when 'postgres'
+          postgres_tls_socket(uri, address)
         else
           raw_tls_socket(uri, address)
         end
       end
 
+      def mysql_tls_socket(uri, address)
+        socket = TCPSocket.new(address, uri.port)
+        length = "#{socket.read(3)}\0".unpack1('L*')
+        _sequence = socket.read(1)
+        body = socket.read(length)
+        initial_handshake_packet = body.unpack('cZ*La8aScSS')
+
+        capabilities = initial_handshake_packet[5] | (initial_handshake_packet[8] << 16)
+
+        ssl_flag = 1 << 11
+        raise 'No TLS support' if (capabilities & ssl_flag).zero?
+
+        socket.write(['2000000185ae7f0000000001210000000000000000000000000000000000000000000000'].pack('H*'))
+        tls_handshake(socket, uri.host)
+      end
+
+      def postgres_tls_socket(uri, address)
+        socket = TCPSocket.new(address, uri.port)
+        socket.write(['0000000804d2162f'].pack('H*'))
+        raise 'Unexpected reply' unless socket.read(1) == 'S'
+
+        tls_handshake(socket, uri.host)
+      end
+
       def smtp_tls_socket(uri, address)
         socket = LineOrientedSocket.new(address, uri.port)
         socket.gets_until_match(/^220 /)