#!/usr/bin/env python3
import hashlib
from Crypto.Cipher import AES
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})}});'

USAGE_STRING = """
Usage:
	signalDesktopDecrypt generate --path <path>
	signalDesktopDecrypt decrypt --path <path>
"""

def decrypt(cipherkey, ciphertext):
	cipher = AES.new(cipherkey, AES.MODE_CBC, iv)
	return Padding.unpad(cipher.decrypt(ciphertext), AES.block_size)

def generate_cipherkey(key, nonce):
	h = HMAC.new(key, digestmod=SHA256)
	h.update(nonce)
	return h.digest()

def generate_key(password):
	m = hashlib.sha256()
	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
		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'))

		# 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))

		# 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()