Make signalDesktopDecrypt configurable

signalDesktopDecrypt.py

@@ -5,17 +5,24 @@ from Crypto.Util import Padding
 from Crypto.Hash import HMAC, SHA256
 import sys
 from getpass import getpass
+from docopt import docopt
+from os import listdir
+
 IV_LENGTH = 16
 NONCE_LENGTH = 16
 MAC_LENGTH = 16
 
+iv = bytearray(IV_LENGTH)
+
 FORMAT_STRING = 'window.Signal.Backup.exportToDirectory(\'{0}\', {{\'key\':new Uint8Array({1})}});'
 
-def read_file(path):
-	with open(path, 'rb') as fh:
-		return fh.read()
+USAGE_STRING = """
+Usage:
+	signalDesktopDecrypt generate --path <path>
+	signalDesktopDecrypt decrypt --path <path>
+"""
 
-def decrypt(cipherkey, iv, ciphertext):
+def decrypt(cipherkey, ciphertext):
 	cipher = AES.new(cipherkey, AES.MODE_CBC, iv)
 	return Padding.unpad(cipher.decrypt(ciphertext), AES.block_size)
 
@@ -29,29 +36,52 @@ def generate_key(password):
 	m.update(password)
 	return m.digest()
 
+def extract_info(filepath):
+	data = ''
+	try:
+		with open(filepath, 'rb') as fh:
+			data = fh.read()
+	except FileNotFoundError:
+		print(filepath + ' was not found, please check --path argument')
+		sys.exit(1)
+
+	# extract none
+	nonce = data[:NONCE_LENGTH]
+	# extract ciphertext
+	ciphertext = data[NONCE_LENGTH:-MAC_LENGTH]
+
+	return nonce, ciphertext
+
 def main():
+	params = docopt(USAGE_STRING)
+
+	if params['generate']:
+		# ask for password
 		password = getpass()
 		key = generate_key(password.encode('utf-8'))
 
 		# generate string to be pasted in Signal-Desktop
-	if len(sys.argv) == 1:
-		print(FORMAT_STRING.format('/tmp/signal/', [elem for elem in key]))
+		print('Use ctrl+i to open the developer tools in the signal desktop app, switch to console and paste the following text:\n')
+		print(FORMAT_STRING.format(params['<path>'], [elem for elem in key]))
+
 		return
+	elif params['decrypt']:
+		# ask for password
+		password = getpass()
+		key = generate_key(password.encode('utf-8'))
 
-	# read encrypted file from disk
-	data = read_file(sys.argv[1])
-
-	# crypto foo
-	iv = bytearray(IV_LENGTH)  # IV is just 16 zeros
-
-	nonce = data[:NONCE_LENGTH]
-	ciphertext = data[NONCE_LENGTH:-MAC_LENGTH]
-
+		# extract messages.zip
+		nonce, ciphertext = extract_info(params['<path>'] + '/messages.zip')
 		cipherkey = generate_cipherkey(key, nonce)
+		with open(params['<path>'] + '/messages.zip', 'wb') as fh:
+			fh.write(decrypt(cipherkey, ciphertext))
 
-	# save decrypted file to 'decrypted'
-	with open('decrypted', 'wb') as fh:
-		fh.write(decrypt(cipherkey, iv, ciphertext))
+		# extract attachments
+		for filename in listdir(params['<path>'] + '/attachments/'):
+			nonce, ciphertext = extract_info(params['<path>'] + '/attachments/' +  filename)
+			cipherkey = generate_cipherkey(key, nonce)
+			with open(params['<path>'] + '/attachments/' + filename, 'wb') as fh:
+				fh.write(decrypt(cipherkey, ciphertext))
 
 if __name__ == '__main__':
 	main()