46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
import glob
|
|
import io
|
|
import traceback
|
|
|
|
from functools import partial
|
|
from multiprocessing.pool import Pool
|
|
from os import path
|
|
|
|
from . import crypt
|
|
|
|
def crypt_dir(dir_in, key, encrypt=False, no_compress=False):
|
|
if not encrypt:
|
|
files = glob.glob(path.join(dir_in, "**/*.utage"), recursive=True)
|
|
else:
|
|
files = glob.glob(path.join(dir_in, "**/*.tsv"), recursive=True)
|
|
|
|
if len(files) == 0:
|
|
raise FileNotFoundError("No valid files found in directory: " + dir_in)
|
|
|
|
pcrypt = partial(crypt_file, key=key, encrypt=encrypt, no_compress=no_compress)
|
|
with Pool() as p:
|
|
p.map(pcrypt, files)
|
|
|
|
def crypt_file(file_in, key, encrypt=False, no_compress=False):
|
|
with open(file_in, "rb") as inf:
|
|
in_data = inf.read()
|
|
try:
|
|
if encrypt:
|
|
if not file_in.endswith(".png") and not file_in.endswith(".tsv"):
|
|
raise ValueError("Invalid File Type for {}".format(file_in))
|
|
if not file_in.endswith(".png") and not no_compress:
|
|
enc_data = crypt.compress(in_data)
|
|
enc_data = crypt.xor_crypt(enc_data, key)
|
|
with io.open(file_in + ".utage", "wb") as output:
|
|
output.write(enc_data)
|
|
else:
|
|
if not file_in.endswith(".utage"):
|
|
raise ValueError("Invalid File Type for {}".format(file_in))
|
|
dec_data = crypt.xor_crypt(in_data, key)
|
|
if not file_in.endswith(".png.utage") and not no_compress:
|
|
dec_data = crypt.decompress(dec_data)
|
|
with io.open(file_in.replace(".utage", ""), "wb") as output:
|
|
output.write(dec_data)
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
print("Error processing file: " + file_in)
|