#! /usr/bin/env python3

# Yandere Lewd Bot, an image posting bot for Pleroma
# Copyright (C) 2022 Anon <yanderefedi@proton.me>
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import sys
import base64
import os
import getpass
from cryptography.fernet import Fernet, InvalidToken
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

from collections import OrderedDict


class EncryptionFail(Exception):
	pass

class PasswordMismatch(Exception):
	pass


# Return string from bytes
def salt_encode(b):
	return base64.urlsafe_b64encode(b).decode()


# Return bytes from string
def salt_decode(s):
	return base64.urlsafe_b64decode(s.encode())


# Ordered Dictionaries
def change_encoding_dict(settings_server, encoding_type):
	return OrderedDict([(k, encoding_type(v)) for k, v in settings_server.items()])


# Return bytes from string
def encode_dict(settings_server):
	return change_encoding_dict(settings_server, str.encode)


# Return string from bytes
def decode_dict(settings_server):
	return change_encoding_dict(settings_server, bytes.decode)


# password: Bytes
# salt: Bytes
def derive_key(password, salt):
	kdf = PBKDF2HMAC(
		algorithm=hashes.SHA256(),
		length=32,
		salt=salt,
		iterations=100000,
		backend=default_backend()
		)
	return base64.urlsafe_b64encode(kdf.derive(password))


# Encryption functions
# message: Bytes
# key: Bytes
def encrypt(message, key):
	f = Fernet(key)
	token = f.encrypt(message)
	return token


# token: Bytes
# key: Bytes
def decrypt(token, key):
	f = Fernet(key)
	message = f.decrypt(token)
	return message


# password: 			bytes()
# salt: 				bytes()
# settings_server: 		dict() -> Byte values
# encryption_function:	encrypt(message, key) : decrypt(token, key):
# Returns settings_server_decrypted dictionary with Byte() values. Will need to use
# ChangeEncodingDict to make them strings (recommended cfg file friendly)
def encrypt_settings(settings_server, password, salt, encryption_function):
	key = derive_key(password, salt)
	settings_server_decrypted = OrderedDict()
	for setting in settings_server:
		settings_server_decrypted[setting] = encryption_function(settings_server[setting], key)
	return settings_server_decrypted

def get_keyfile(keyfile=None):
	if keyfile is not None:
		with open(keyfile, "rb") as f:
			return f.read()
	return None


def get_pass(q):
	try:
		return getpass.getpass(q).encode()
	except KeyboardInterrupt:
		raise EncryptionFail("\nQuitting...")


def settings_server_encrypt(settings_server, keyfile=None):
	try:
		settings_server = encode_dict(settings_server)
		salt = os.urandom(16)
		password = get_keyfile(keyfile)
		if password is None:
			password = get_pass("Enter password: ")
			password2 = get_pass("Enter password: ")
			if password != password2:
				raise PasswordMismatch("Passwords do not match")
		encrypted = encrypt_settings(settings_server, password, salt, encrypt)
		return salt_encode(salt), decode_dict(encrypted)
	except PasswordMismatch as e:
		raise EncryptionFail(str(e))
	except Exception as e:
		err = str(e)
		raise EncryptionFail("Encrypt Error: {}".format(err))


def settings_server_decrypt(settings_server, settings_encrypt, keyfile=None):
	try:
		if not settings_encrypt["encrypt"]:
			return settings_server
		settings_server = encode_dict(settings_server)
		password = get_keyfile(keyfile or settings_encrypt["keyfile"]) or get_pass("Enter password: ")
		salt = salt_decode(settings_encrypt["salt"])
		decrypted = encrypt_settings(settings_server, password, salt, decrypt)
		return decode_dict(decrypted)
	except base64.binascii.Error:
		raise EncryptionFail("Salt is invalid")
	except InvalidToken:
		raise EncryptionFail("Password or token is incorrect")
	except Exception as e:
		err = str(e)
		raise EncryptionFail("Decrypt Error: {}".format(err))


def main():
	import argparse
	from pprint import pformat

	default_cfg = "cfg"

	parser = argparse.ArgumentParser(
									description="A class to encrypt server credentials",
									epilog="There are no additional parameters.",
									add_help=True	)
	parser.add_argument("--encrypt", help="Generate encrypted authentication.", action="store_true")
	parser.add_argument("--decrypt", help="Decrypt encrypted authentication", action="store_true")
	parser.add_argument("--recrypt", help="Recrypt encrypted authentication", action="store_true")
	parser.add_argument("-k", "--keyfile", help="Keyfile used for decryption", default=None)
	parser.add_argument("-c", "--cfg", help="Specify config file.", default=default_cfg)

	arguments = parser.parse_args()

	if arguments.recrypt:
		arguments.encrypt = True
		arguments.decrypt = True

	if arguments.encrypt or arguments.decrypt:
		import importlib
		cfg = importlib.import_module(arguments.cfg)
		settings_server = cfg.settings_server
		settings_encrypt = cfg.settings_encrypt
		keyfile = arguments.keyfile or settings_encrypt["keyfile"]

		if arguments.decrypt and arguments.encrypt:
			print("Re-encrypting")

		if arguments.decrypt:		# arguments.decrypt
			print("Decrypt...")
			settings_server = settings_server_decrypt(settings_server, settings_encrypt, arguments.keyfile)
			settings_encrypt = OrderedDict([
				("encrypt",	False),
				("salt",	settings_encrypt["salt"]),
				("keyfile", arguments.keyfile)
			])

		if arguments.encrypt:
			print("Encrypt...")
			salt, settings_server = settings_server_encrypt(settings_server, keyfile)
			settings_encrypt = OrderedDict([
				("encrypt",	True),
				("salt",	salt),
				("keyfile", arguments.keyfile)
			])

		print("settings_server = {}".format(pformat(settings_server)))
		print("settings_encrypt = {}".format(pformat(settings_encrypt)))
	return 0


if __name__ == "__main__":
	try:
		sys.exit(main())
	except EncryptionFail as e:
		print(e)
		sys.exit(1)