#!/usr/bin/env python3
import hashlib
from Crypto.Cipher import AES
from Crypto.Util import Padding
from Crypto.Hash import HMAC, SHA256
import sys
IV_LENGTH = 16
NONCE_LENGTH = 16
MAC_LENGTH = 16

FORMAT_STRING = 'window.Signal.Backup.exportToDirectory(\'{0}\', {{\'key\':new Uint8Array({1})}});'

def read_file(path):
	with open(path, 'rb') as fh:
		return fh.read()

def decrypt(cipherkey, iv, 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 main():
	password = input('Password: ')
	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]))
		return

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

	cipherkey = generate_cipherkey(key, nonce)

	# save decrypted file to 'decrypted'
	with open('decrypted', 'wb') as fh:
		fh.write(decrypt(cipherkey, iv, ciphertext))

if __name__ == '__main__':
	main()