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
 from OpenSSL import SSL
 import base64, pickle
 import getopt, sys, os, time

+import io

 
 verbose = False
 
 
 verbose = False
 

@@ -122,6 +123,10 @@ class ProxyFactory(protocol.Factory):
         self.authtoken = authtoken
         self.machine = machine
 
         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:
 def main():
     global verbose
     try:

@@ -175,7 +180,7 @@ def main():
 
     # Unpack authentication token
     try:
 
     # 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"]
         machine = token_inner["machine"]
         connect_host = token_inner["connect_host"]
         connect_port = token_inner["connect_port"]