#   scipy-yli: Helpful SciPy utilities and recipes
#   Copyright © 2022  Lee Yingtong Li (RunasSudo)
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU Affero 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 Affero General Public License for more details.
#
#   You should have received a copy of the GNU Affero General Public License
#   along with this program.  If not, see <https://www.gnu.org/licenses/>.

from Crypto.Cipher import AES
from Crypto.Protocol.KDF import scrypt
from Crypto.Random import get_random_bytes

import getpass
import io
import lzma
import pickle

def pickle_write_encrypted(df, fname, salt=None):
	"""
	Write the DataFrame to an encrypted file
	
	DataFrame is serialised with Pickle, compressed with LZMA, then encrypted
	Encryption is AES-256-CTR
	Password is derived using scrypt KDF
	"""
	
	# Get random salt
	if salt is None:
		salt = get_random_bytes(16)
	
	# Serialise Pandas and compress
	pickle_pt = io.BytesIO()
	with lzma.LZMAFile(pickle_pt, 'wb') as f:
		pickle.dump(df, f)
	
	# Generate key using scrypt
	password = getpass.getpass('Password: ')
	key = scrypt(password, salt, 256//8, 2**20, 8, 1)
	
	# Encrypt Pandas plaintext
	cipher = AES.new(key, AES.MODE_CTR)
	ct_bytes = cipher.encrypt(pickle_pt.getvalue())
	
	# Write to file
	with open(fname, 'wb') as f:
		pickle.dump((salt, cipher.nonce, ct_bytes), f)

def pickle_write_compressed(df, fname):
	"""
	Write the DataFrame to a compressed file
	
	DataFrame is serialised with Pickle and compressed with LZMA
	"""
	
	# Serialise Pandas and compress
	with open(fname, 'wb') as f:
		with lzma.LZMAFile(f, 'wb') as lf:
			pickle.dump(df, lf)

def pickle_read_encrypted(fname):
	"""
	Read a DataFrame from an encrypted file
	"""
	
	# Read ciphertext data
	with open(fname, 'rb') as f:
		salt, nonce, ct_bytes = pickle.load(f)
	
	# Generate key using scrypt
	password = getpass.getpass('Password: ')
	key = scrypt(password, salt, 256//8, 2**20, 8, 1)
	
	# Decrypt to Pandas plaintext
	cipher = AES.new(key, AES.MODE_CTR, nonce=nonce)
	pickle_pt = io.BytesIO()
	pickle_pt.write(cipher.decrypt(ct_bytes))
	pickle_pt.seek(0)
	
	# Uncompress and deserialise to Pandas
	return pickle.load(lzma.LZMAFile(pickle_pt, 'rb'))

def pickle_read_compressed(fname):
	"""
	Read a DataFrame from an compressed (but not encrypted) file
	"""
	
	# Uncompress and deserialise to Pandas
	with open(fname, 'rb') as f:
		return pickle.load(lzma.LZMAFile(f, 'rb'))