From 49b7842f47ac7f62560f7f2e17ea345f32762c43 Mon Sep 17 00:00:00 2001
From: Sharat M R <cosmosgenius@gmail.com>
Date: Wed, 18 Mar 2020 17:16:47 +0530
Subject: [PATCH] Add password support for vnc

---
 Preferences.plist | 12 ++++++++++++
 Tweak.xm          | 40 +++++++++++++++++++++++++++++++++++++---
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/Preferences.plist b/Preferences.plist
index 94f8d0b..0037786 100644
--- a/Preferences.plist
+++ b/Preferences.plist
@@ -33,6 +33,18 @@
             <key>PostNotification</key>
             <string>com.cosmosgenius.screendump/preferences.changed</string>
         </dict>
+        <dict>
+            <key>cell</key>
+            <string>PSSecureEditTextCell</string>
+            <key>defaults</key>
+            <string>com.cosmosgenius.screendump</string>
+            <key>key</key>
+            <string>CCSPassword</string>
+            <key>label</key>
+            <string>Password</string>
+            <key>PostNotification</key>
+            <string>com.cosmosgenius.screendump/preferences.changed</string>
+        </dict>
     </array>
 </dict>
 </plist>
diff --git a/Tweak.xm b/Tweak.xm
index 58e6fc3..1c5d04c 100644
--- a/Tweak.xm
+++ b/Tweak.xm
@@ -7,6 +7,7 @@
 #define kSettingsPath [NSHomeDirectory() stringByAppendingPathComponent:@"/Library/Preferences/com.cosmosgenius.screendump.plist"]
 
 static bool CCSisEnabled = true;
+static NSString *CCSPassword = nil;
 static rfbScreenInfoPtr screen;
 static bool isVNCRunning;
 // static NSCondition *condition_;
@@ -42,14 +43,30 @@ extern "C" kern_return_t IOMobileFramebufferSwapSetLayer(
 extern "C" void IOSurfaceFlushProcessorCaches(IOSurfaceRef buffer);
 extern "C" int IOSurfaceLock(IOSurfaceRef surface, uint32_t options, uint32_t *seed);
 extern "C" int IOSurfaceUnlock(IOSurfaceRef surface, uint32_t options, uint32_t *seed);
-static void VNCUpdateRunState(bool shouldStart);
 
 static IOSurfaceAcceleratorRef accelerator;
 static IOSurfaceRef static_buffer;
 
+static void VNCSettings(bool shouldStart, NSString *password);
+static void VNCUpdateRunState(bool shouldStart);
 static void handleVNCKeyboard(rfbBool down, rfbKeySym key, rfbClientPtr client);
 static void handleVNCPointer(int buttons, int x, int y, rfbClientPtr client);
 
+static rfbBool VNCCheck(rfbClientPtr client, const char *data, int size) {
+    NSString *password = reinterpret_cast<NSString *>(screen->authPasswdData);
+    if(!password) {
+        return TRUE;
+    }
+    if ([password length] == 0) {
+        return TRUE;
+    }
+    NSAutoreleasePool *pool([[NSAutoreleasePool alloc] init]);
+    rfbEncryptBytes(client->authChallenge, const_cast<char *>([password UTF8String]));
+    bool good(memcmp(client->authChallenge, data, size) == 0);
+    [pool release];
+    return good;
+}
+
 static void VNCSetup() {
     int argc(1);
     char *arg0(strdup("ScreenDumpVNC"));
@@ -61,6 +78,7 @@ static void VNCSetup() {
     screen->serverFormat.blueShift = bits_per_sample * 0;
     screen->kbdAddEvent = &handleVNCKeyboard;
     screen->ptrAddEvent = &handleVNCPointer;
+    screen->passwordCheck = &VNCCheck;
     free(arg0);
     VNCUpdateRunState(CCSisEnabled);
 }
@@ -116,10 +134,25 @@ static void initialBuffer() {
     }
 }
 
+static void VNCSettings(bool shouldStart, NSString* password) {
+    CCSisEnabled = shouldStart;
+    if(password) {
+        CCSPassword = password;
+    }
+    NSString *sEnabled = CCSisEnabled ? @"YES": @"NO";
+    // NSLog(@"screendump: Settings(Enabled:%@, Password:%@)", sEnabled, CCSPassword);
+    VNCUpdateRunState(CCSisEnabled);
+}
+
 static void VNCUpdateRunState(bool shouldStart) {
     if(screen == NULL) {
         return;
     }
+    if(CCSPassword && CCSPassword.length) {
+        screen->authPasswdData = (void *) CCSPassword;
+    } else {
+        screen->authPasswdData = NULL;
+    }
     if(shouldStart == isVNCRunning) {
         return;
     }
@@ -185,10 +218,11 @@ static void loadPrefs(void)
     }
 
     if(prefs) {
-        CCSisEnabled = [prefs objectForKey:@"CCSisEnabled"] ? [[prefs objectForKey:@"CCSisEnabled"] boolValue] : CCSisEnabled;
+        bool isEnabled = [prefs objectForKey:@"CCSisEnabled"] ? [[prefs objectForKey:@"CCSisEnabled"] boolValue] : CCSisEnabled;
+        NSString *password = [prefs objectForKey:@"CCSPassword"];
         [prefs release];
+        VNCSettings(isEnabled, password);
     }
-    VNCUpdateRunState(CCSisEnabled);
 }
 
 %hookf(kern_return_t, IOMobileFramebufferSwapSetLayer, IOMobileFramebufferRef fb, int layer, IOSurfaceRef buffer, CGRect bounds, CGRect frame, int flags) {