From edc149af29242178091b2d6fcd42c3ef0851644b Mon Sep 17 00:00:00 2001
From: Marios Levogiannis <mlevogiannis@users.noreply.github.com>
Date: Fri, 29 Jul 2022 18:40:09 +0300
Subject: [PATCH] Fix CVE-2022-36436 - Authentication bypass in RFB security
 handshake (#1)

---
 vncap/tests/test_protocol.py | 16 +++++++++++++++-
 vncap/vnc/protocol.py        |  7 ++-----
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/vncap/tests/test_protocol.py b/vncap/tests/test_protocol.py
index 151d059..81bef84 100644
--- a/vncap/tests/test_protocol.py
+++ b/vncap/tests/test_protocol.py
@@ -13,6 +13,9 @@ def write(self, data):
     def loseConnection(self):
         self.lost = True
 
+    def pauseProducing(self):
+        pass
+
 class TestVNCServerAuthenticator(unittest.TestCase):
 
     def setUp(self):
@@ -29,9 +32,20 @@ def test_connectionMade(self):
     def test_check_version(self):
         self.t.buf = ""
         self.p.check_version("RFB 003.008\n")
-        self.assertEqual(self.t.buf, "\x02\x01\x02")
+        self.assertEqual(self.t.buf, "\x01\x02")
 
     def test_check_invalid_version(self):
         self.t.buf = ""
         self.p.check_version("RFB 002.000\n")
         self.assertTrue(self.t.lost)
+
+    def test_select_security_type_none(self):
+        self.t.buf = ""
+        self.p.select_security_type("\x01")
+        self.assertTrue(self.t.lost)
+
+    def test_select_security_type_vnc_auth(self):
+        self.t.buf = ""
+        self.p.select_security_type("\x02")
+        self.assertFalse(self.t.lost)
+        self.assertEqual(len(self.t.buf), 16)
diff --git a/vncap/vnc/protocol.py b/vncap/vnc/protocol.py
index aee50d3..d1351bd 100644
--- a/vncap/vnc/protocol.py
+++ b/vncap/vnc/protocol.py
@@ -73,8 +73,8 @@ def check_version(self, version):
 
         if version == self.VERSION:
             log.msg("Client version %s is valid" % version.strip())
-            # Hardcoded: 2 security types: None and VNC Auth.
-            self.transport.write("\x02\x01\x02")
+            # Hardcoded: 1 security type: VNC Auth.
+            self.transport.write("\x01\x02")
             return self.select_security_type, 1
         else:
             log.err("Can't handle VNC version %r" % version)
@@ -93,9 +93,6 @@ def select_security_type(self, security_type):
             self.transport.write(self.challenge)
 
             return self.vnc_authentication_result, 16
-        elif security_type == 1:
-            # No authentication. Just move to the SecurityResult.
-            self.authenticated()
         else:
             log.err("Couldn't agree on an authentication scheme!")
             self.transport.loseConnection()