import sys
import zlib
from argparse import ArgumentParser
from io import BytesIO
from zipfile import ZipFile
from os.path import basename as get_basename

KEY = b"PUuhgrUGGIUh9JHGUIGIUHGokfewrofijU"
DECRYPTED_DEX_PATH = r"decrypted_{}.dex"

ZLIB_HEADER = b'x\x9c'


def decrypt_file(key, f, filename):
	# Read encrypted DEX
	output_buffer = bytearray()
	compressed_dex = f.read()
	print(f"Dex length: {len(compressed_dex)}")
	if not len(compressed_dex):
		return None

	# First stage ZLib Decompression
	try:
		decompressed = zlib.decompress(compressed_dex)
	except zlib.error as e:
		raise Exception(f"Could not decompress first stage: {e}")
	decompressed_dex_file = BytesIO(decompressed)

	# XOR Decryption
	while True:
		data = decompressed_dex_file.read(8192)
		if not data:
			break
		sys.stdout.write(".")
		sys.stdout.flush()
		for bytes_changed, b in enumerate(data):
			transform0 = [(key[9] << 16) | key[8], (key[11] << 16) | key[10]]
			transform1 = transform0[int(bytes_changed % 8 / 4)]
			transform2 = (transform1 >> ((bytes_changed % 4) << 3)) & 0x000000FF
			transform3 = (transform2 ^ data[bytes_changed]) & 0x000000FF
			output_buffer.append(transform3)
	sys.stdout.write("\n")

	# Second stage ZLib Decompression
	try:
		result = zlib.decompress(output_buffer)
	except zlib.error as e:
		raise Exception(f"Could not decompress second stage: {e}")

	# Output DEX file
	outfile = DECRYPTED_DEX_PATH.format(filename)
	print(f"Writing to: {outfile}")
	with open(outfile, "wb") as f:
		f.write(result)
	print("Success!\n")


def main():
	# Parse arguments
	parser = ArgumentParser(description='Flubot Dex Extractor')
	parser.add_argument('-f', '--apk', help='Flubot APK', required=True)
	parser.add_argument('-k', '--key', help='Decryption key', default=b"PUuhgrUGGIUh9JHGUIGIUHGokfewrofijU")
	args = vars(parser.parse_args())

	# Open APK and find
	with ZipFile(args.get('apk')) as zf:
		potential_files = [x for x in zf.filelist if x.file_size > 900000 and x.filename != "classes.dex"]
		if not len(potential_files):
			raise Exception("No file candidates was found inside the APK")

		for zipfile_info in potential_files:
			with zf.open(zipfile_info) as zipentry:
				if zipentry.read(2) != ZLIB_HEADER:
					continue
				zipentry.seek(0)
				filename = get_basename(zipfile_info.filename)
				try:
					decrypt_file(args.get('key'), zipentry, filename)
				except Exception as e:
					print(f"[Error]: {filename} - {e}")


if __name__ == "__main__":
	try:
		main()
	except Exception as e:
		print(f"[Error]: {e}")