126 lines
3.3 KiB
Python
126 lines
3.3 KiB
Python
# 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/>.
|
|
|
|
import getpass
|
|
import io
|
|
import lzma
|
|
import pickle
|
|
|
|
def pickle_write_encrypted(df, fname):
|
|
"""
|
|
Write the DataFrame to an encrypted file
|
|
|
|
Prompts the user for a password.
|
|
|
|
The DataFrame is serialised with *pickle*, compressed with LZMA, then encrypted.
|
|
Encryption uses AES-256-CTR.
|
|
Encryption key is derived from the password using the scrypt KDF.
|
|
|
|
:param df: Data to serialise
|
|
:type df: DataFrame
|
|
:param fname: Filename to write data to
|
|
:type fname: str
|
|
"""
|
|
|
|
from Crypto.Cipher import AES
|
|
from Crypto.Protocol.KDF import scrypt
|
|
from Crypto.Random import get_random_bytes
|
|
|
|
# Get random salt
|
|
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
|
|
|
|
The DataFrame is serialised with *pickle*, then compressed with LZMA.
|
|
|
|
:param df: Data to serialise
|
|
:type df: DataFrame
|
|
:param fname: Filename to write data to
|
|
:type fname: str
|
|
"""
|
|
|
|
# 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
|
|
|
|
See :meth:`yli.pickle_write_encrypted`.
|
|
|
|
:param fname: Filename to read data from
|
|
:type fname: str
|
|
|
|
:rtype: DataFrame
|
|
"""
|
|
|
|
from Crypto.Cipher import AES
|
|
from Crypto.Protocol.KDF import scrypt
|
|
|
|
# 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 a compressed file
|
|
|
|
See :meth:`yli.pickle_write_compressed`.
|
|
|
|
:param fname: Filename to read data from
|
|
:type fname: str
|
|
|
|
:rtype: DataFrame
|
|
"""
|
|
|
|
# Uncompress and deserialise to Pandas
|
|
with open(fname, 'rb') as f:
|
|
return pickle.load(lzma.LZMAFile(f, 'rb'))
|