Subclass pickle.Unpickler for security

Signed-off-by: Anders Kaseorg <andersk@mit.edu>

diff --git a/invirt-vnc-client b/invirt-vnc-client
index b789c56..f71366a 100755 (executable)
--- a/invirt-vnc-client
+++ b/invirt-vnc-client
@@ -3,6 +3,7 @@ from twisted.internet import reactor, ssl, protocol, error
 from OpenSSL import SSL
 import base64, pickle
 import getopt, sys, os, time
+import io
 
 verbose = False
 
@@ -122,6 +123,10 @@ class ProxyFactory(protocol.Factory):
         self.authtoken = authtoken
         self.machine = machine
 
+class SafeUnpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        raise pickle.UnpicklingError("globals are forbidden")
+
 def main():
     global verbose
     try:
@@ -175,7 +180,7 @@ def main():
 
     # Unpack authentication token
     try:
-        token_inner = pickle.loads(base64.urlsafe_b64decode((authtoken.split("."))[0]))
+        token_inner = SafeUnpickler(io.BytesIO(base64.urlsafe_b64decode((authtoken.split("."))[0]))).load()
         machine = token_inner["machine"]
         connect_host = token_inner["connect_host"]
         connect_port = token_inner["connect_port"]