Fixed minor bugs in decrypting project

Project now decrypts properly and fails out loudly when incorrect
password is supplied. Must supply project name via command line now.

pappyproxy/crypto.py

@@ -17,7 +17,7 @@ from twisted.internet import reactor, defer
 class Crypto(object):
     def __init__(self, sessconfig):
         self.config = sessconfig
-        self.archive = self.config.archive 
+        self.archive = sessconfig.archive 
         self.compressor = compress.Compress(sessconfig)
         self.key = None
         self.password = None
@@ -31,15 +31,15 @@ class Crypto(object):
         # Leave the crypto working directory
         os.chdir('../')
 
-        # Get the password and salt, then derive the key
-        self.crypto_ramp_up()
-        
         self.compressor.compress_project()
          
         # Create project and crypto archive 
         archive_file = open(self.archive, 'rb') 
         archive_crypt = open(self.config.crypt_file, 'wb')
          
+        # Get the password and salt, then derive the key
+        self.crypto_ramp_up()
+        
         # Encrypt the archive read as a bytestring
         fern = Fernet(self.key)
         crypt_token = fern.encrypt(archive_file.read())
@@ -61,29 +61,39 @@ class Crypto(object):
         """
 
         # If project hasn't been encrypted before, setup crypt working directory
-        crypt_fp = os.path.join(os.getcwd(), self.config.crypt_file)
+        if not os.path.isfile(self.config.crypt_file):
-        if not os.path.isfile(crypt_fp):
             os.mkdir(self.config.crypt_dir)
 
             project_files = self.config.get_project_files()
             for pf in project_files:
                 shutil.copy2(pf, self.config.crypt_dir)
             os.chdir(self.config.crypt_dir)
+            return True
         
         # Otherwise, decrypt and decompress the project 
         else: 
-            self.crypto_ramp_up()
-            fern = Fernet(self.key)
-
-            # Decrypt the project archive
             archive_crypt = open(self.config.crypt_file, 'rb').read()
             archive_file = open(self.config.archive, 'wb')
-            try:
+
-                archive = fern.decrypt(archive_crypt)
+            retries = 3
-            except InvalidToken:
+            while True:
-                raise PappyException("Problem decrypting the file, restart pappy to try again")
+                try:
-                reactor.stop()
+               	    self.crypto_ramp_up()
-                defer.returnValue(None)
+               	    fern = Fernet(self.key)
+                    archive = fern.decrypt(archive_crypt)
+                    break
+                except InvalidToken:
+                    print "Invalid password"
+                    retries -= 1
+                    # Quit pappy if user doesn't retry
+                    # or if all retries exhuasted
+                    if not self.confirm_password_retry() or retries <= 0:
+                        return False
+                    else:
+                        self.password = None
+                        self.key = None
+                        self.salt = None
+                        pass 
                         
             archive_file.write(archive)
             archive_file.close()
@@ -94,6 +104,14 @@ class Crypto(object):
             self.delete_crypt_files()
             
             os.chdir(self.config.crypt_dir)
+            return True
+            
+    def confirm_password_retry(self):
+        answer = raw_input("Would you like to re-enter your password? (y/n)").strip()
+        if answer[0] == "y" or answer[0] == "Y":
+            return True
+        else:
+            return False
             
     def crypto_ramp_up(self):
         if not self.password:
@@ -108,7 +126,7 @@ class Crypto(object):
         """
         encoded_passwd = ""
         try:
-            passwd = raw_input("Enter a password: ")
+            passwd = raw_input("Enter a password: ").strip()
             self.password = passwd.encode("utf-8")
         except:
             raise PappyException("Invalid password, try again")

pappyproxy/pappy.py

@@ -65,7 +65,6 @@ class PappySession(object):
         self.delete_data_on_quit = False
         self.ports = None
         self.crypto = crypto.Crypto(sessconfig)
-        self.password = None
 
     @defer.inlineCallbacks
     def start(self):
@@ -155,7 +154,13 @@ class PappySession(object):
 
     @defer.inlineCallbacks
     def decrypt(self):
-        yield self.crypto.decrypt_project()
+        # Attempt to decrypt project archive
+        if self.crypto.decrypt_project():
+            yield True
+        # Quit pappy on failure
+        else:
+            reactor.stop() 
+            defer.returnValue(None)
              
     @defer.inlineCallbacks
     def cleanup(self, ignored=None):
@@ -175,7 +180,12 @@ def parse_args():
 
     parser = argparse.ArgumentParser(description='An intercepting proxy for testing web applications.')
     parser.add_argument('-l', '--lite', help='Run the proxy in "lite" mode', action='store_true')
-    parser.add_argument('-c', '--crypt', type=str, nargs='?', help='Start pappy in "crypto" mode, optionally supply a name for the encrypted project archive [CRYPT]')
+    try:
+        parser.add_argument('-c', '--crypt', type=str, nargs=1, help='Start pappy in "crypto" mode, must supply a name for the encrypted project archive [CRYPT]')
+    except:
+        print 'Must supply a project name: pappy -c <project_name>'
+        reactor.stop() 
+        defer.returnValue(None)
 
     args = parser.parse_args(sys.argv[1:])
     settings = {}
@@ -186,9 +196,8 @@ def parse_args():
         settings['lite'] = False
 
     if args.crypt:
-        settings['crypt'] = args.crypt 
+        # Convert from single-item list produced by argparse `nargs=1`
-    elif args.crypt == "":
+        settings['crypt'] = args.crypt[0].encode('utf-8') 
-        settings['crypt'] = 'project.crypt'
     else:
         settings['crypt'] = None