forked from carlospuenteg/File-Injector
-
Notifications
You must be signed in to change notification settings - Fork 0
/
cryptography.py
98 lines (68 loc) · 2.84 KB
/
cryptography.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
import os
import re
from pathlib import Path
from cryptography.fernet import Fernet, InvalidToken
from constants.constants import ENCRYPTION_KEYS_PATH
from classes.Options import Options
from utils.ctxt import *
from utils.input import get_path, get_valid_filename
# Generate a key with Fernet.generate_key(), save it to a file and return the path
def gen_key(new_filename:str=None) -> str:
key = Fernet.generate_key()
if new_filename:
stem = Path(new_filename).stem
path = f"{ENCRYPTION_KEYS_PATH}/{stem}.key"
else:
idxs = [0]
for file in os.listdir(ENCRYPTION_KEYS_PATH):
if re.match(r"key\d+\.key", file):
idxs.append(int(file[3:-4]))
new_idx = max(idxs) + 1
path = f"{ENCRYPTION_KEYS_PATH}/key{new_idx}.key"
with open(path, 'wb') as f:
f.write(key)
print(f"Key generated and saved to {ctxt(path, Fore.GREEN)}")
return path
# Load the key from the path
def load_key(path:str) -> bytes:
stem = Path(path).stem
with open(path, 'rb') as f:
key = f.read()
return key
# Choose or generate a key from a menu, validate it and return the Fernet object
def get_fernet() -> Fernet:
print("\nDo you want to use an existing key or generate a new one?")
option = Options(["Existing key", "New key"]).get_choice()
if option == 0:
key_path = get_path([ENCRYPTION_KEYS_PATH], "Filename of the key file: ", [".key"])
if option == 1:
key_filename = get_valid_filename("Filename of the new key file (blank for default): ", [".key"], allow_blank=True)
key_path = gen_key(key_filename)
return Fernet(load_key(key_path))
# Choose the path of the key file and return the Fernet object
def choose_fernet() -> Fernet:
key_path = get_path([ENCRYPTION_KEYS_PATH], "\nFilename of the key file: ", [".key"])
return Fernet(load_key(key_path))
# Return a list of all the paths of the key files
def get_key_paths() -> list:
keys_paths = []
for file in os.listdir(ENCRYPTION_KEYS_PATH):
if file.endswith(".key"):
keys_paths.append(f"{ENCRYPTION_KEYS_PATH}/{file}")
return keys_paths
# Decrypt the content of the file and filename trying with all the key files
def decrypt_content(file:bytes, filename:bytes) -> dict:
for key_path in get_key_paths():
fernet = Fernet(load_key(key_path))
try:
file = fernet.decrypt(file)
filename = fernet.decrypt(filename)
print(ctxt(f'Decrypted with "{key_path}"', Fore.GREEN))
return {
"file": file,
"filename": filename
}
except InvalidToken:
continue
# If the file couldn't be decrypted with any key, raise an Exception
raise Exception(f'Decryption key for the file not found in "{ENCRYPTION_KEYS_PATH}"')