-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
358 lines (296 loc) · 14.5 KB
/
main.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
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
import os
import sys
import base64
from Cryptodome.Hash import SHA512
from Cryptodome.PublicKey import RSA
from Cryptodome.Random import get_random_bytes
from Cryptodome.Cipher import AES, PKCS1_OAEP
from Cryptodome.Protocol.KDF import PBKDF2
class Program:
def __init__(self):
self.session_derived_key = None
self.public_key = None
self.private_key = None
self.root = None
def derive_session_key(self, key: str, salt: bytes) -> bytes:
"""
Derives a session key using PBKDF2.
:param key : The base key to derive from.
:param salt : Salt for the key derivation.
Returns:
bytes: The derived session key.
"""
if self.session_derived_key is None:
self.session_derived_key = PBKDF2(key, salt, 32, count=1000000, hmac_hash_module=SHA512)
return self.session_derived_key
def clear_session_key(self) -> None:
"""
Clears the global session key.
"""
self.session_derived_key = None
def generate_keys(self) -> None:
"""
Generates RSA public and private keys.
Returns:
tuple[bytes, bytes]: A tuple containing (public_key, private_key).
"""
key = RSA.generate(2048)
self.public_key = key.public_key().export_key()
self.private_key = key.export_key()
@staticmethod
def save_key(key: bytes, filename: str, password: str = None, is_name_key: bool = False,
salt: bytes = None) -> None:
"""
This function save the keys into a file for security. Can be encrypted with a passphrase.
:param key: The key to save.
:param filename: The filename to save the key to.
:param password: Password to encrypt the key. Defaults to None. Optional.
:param is_name_key: Whether this is a name key. Defaults to False. Optional.
:param salt: Salt for name key encryption. Required if is_name_key is True. Optional Depending.
"""
if is_name_key:
with open(filename, "wb") as file:
file.write(salt)
file.write(key)
else:
if password:
key = RSA.import_key(key).export_key(passphrase=password, pkcs=8,
protection="scryptAndAES128-CBC")
with open(filename, "wb") as file:
file.write(key)
@staticmethod
def load_key(filename: str, password: str = None, is_name_key: bool = False):
"""
Loads a key from a file.
:param filename: The filename to load the key from.
:param password: Password to decrypt the key. Defaults to None. Optional.
:param is_name_key: Whether this is a name key. Defaults to False. Optional.
"""
with open(filename, "rb") as file:
if is_name_key:
salt = file.read(16)
key = file.read()
return salt, key
else:
encoded_key = file.read()
if password:
key = RSA.import_key(encoded_key, passphrase=password)
else:
key = RSA.import_key(encoded_key)
return key.export_key()
def encrypt_dir(self, directory: str, public_key: bytes, name_key: str, salt: bytes) -> str:
"""
Encrypts an entire directory.
:param directory: The location of where your files are located at.
:param public_key: The public key generated using RSA.
:param name_key : The private key used to encrypt the file names.
:param salt: Salt for name encryption.
WARNING:
YOU MIGHT RISK RUINING YOUR ENTIRE SYSTEM AND CAUSING IRRECOVERABLE FILE LOSS.
ONLY RUN THIS FUNCTION WHEN YOU ARE ABSOLUTELY SURE.
"""
# Creates aes key
aes_key = get_random_bytes(16)
rsa_key = RSA.import_key(public_key)
cipher_rsa = PKCS1_OAEP.new(rsa_key)
enc_aes_key = cipher_rsa.encrypt(aes_key)
# Recursively browse directories
for root, dirs, files in os.walk(directory, topdown=False):
for file_name in files:
file_path = os.path.join(root, file_name)
self.encrypt_file(file_path, aes_key, enc_aes_key, name_key, salt) # Perform file encryption
for dir_name in dirs:
dir_path = os.path.join(root, dir_name)
name_change = self.encrypt_name(dir_name, name_key, salt)
encrypted_path = os.path.join(root, name_change)
os.rename(dir_path, encrypted_path)
print(f"Folder {dir_path} has been encrypted and renamed to {encrypted_path}")
# Perform filename encryption
# Perform root name encryption
root_name = os.path.basename(directory)
root_change = self.encrypt_name(root_name, name_key, salt)
encrypted_root = os.path.join(os.path.dirname(directory), root_change)
os.rename(directory, encrypted_root)
print(f"Root {directory} has been encrypted and renamed to {encrypted_root}")
self.clear_session_key()
return encrypted_root
def decrypt_dir(self, encrypted_dir_path: str, private_key: bytes, name_key: str, salt: bytes) -> str:
"""
Decrypts an entire encrypted directory.
:param encrypted_dir_path: The path of the encrypted directory.
:param private_key: The private key for decryption.
:param name_key: The key used to decrypt file and directory names.
:param salt: Salt for name decryption.
Returns:
str: The path of the decrypted directory.
"""
decrypted_root_name = self.decrypt_name(os.path.basename(encrypted_dir_path), name_key, salt)
decrypted_root_path = os.path.join(os.path.dirname(encrypted_dir_path), decrypted_root_name)
os.rename(encrypted_dir_path, decrypted_root_path)
for root, dirs, files in os.walk(decrypted_root_path, topdown=False):
for enc_file_name in files:
enc_file_path = os.path.join(root, enc_file_name)
self.decrypt_file(enc_file_path, private_key, name_key, salt)
for enc_dir_name in dirs:
enc_dir_path = os.path.join(root, enc_dir_name)
decrypted_dir_name = self.decrypt_name(enc_dir_name, name_key, salt)
decrypted_path = os.path.join(root, decrypted_dir_name)
os.rename(enc_dir_path, decrypted_path)
print(f"Directory '{decrypted_root_path}' has been decrypted.")
self.clear_session_key()
return decrypted_root_path
def encrypt_file(self, file_directory: str, aes_key: bytes, enc_aes_key: bytes, name_key: str, salt: bytes) -> None:
"""
Encrypts a single file.
:param file_directory: The path of the file to encrypt.
:param aes_key: The AES key for file content encryption.
:param enc_aes_key: The encrypted AES key.
:param name_key: The key for encrypting the file name.
:param salt: Salt for name encryption.
WARNING:
THIS IS THE ACTUAL ENCRYPTION FUNCTION. DO NOT USE THIS UNLESS YOU ARE ABSOLUTELY SURE THAT YOU KNOW WHAT YOU ARE
DOING. THIS MAY POTENTIALLY LEAD TO IRRECOVERABLE FILE, DATA LOSS AND POTENTIALLY RUIN YOUR ENTIRE SYSTEM.
"""
with open(file_directory, "rb") as files:
data = files.read()
cipher_aes = AES.new(aes_key, AES.MODE_GCM)
ciphertext, tag = cipher_aes.encrypt_and_digest(data)
encrypted_name = self.encrypt_name(os.path.basename(file_directory), name_key, salt)
encrypted_path = os.path.join(os.path.dirname(file_directory), encrypted_name)
with open(encrypted_path, "wb") as f:
f.write(enc_aes_key)
f.write(cipher_aes.nonce)
f.write(tag)
f.write(ciphertext)
os.remove(file_directory)
print(f"File '{file_directory}' has been encrypted and renamed to '{encrypted_path}'")
def decrypt_file(self, encrypted_file_path: str, private_key: bytes, name_key: str, salt: bytes) -> None:
"""
Decrypts a single encrypted file.
:param encrypted_file_path: The path of the encrypted file.
:param private_key: The private key for decryption.
:param name_key: The key for decrypting the file name.
:param salt: Salt for name decryption.
"""
rsa_key = RSA.import_key(private_key)
cipher_rsa = PKCS1_OAEP.new(rsa_key)
with open(encrypted_file_path, "rb") as f:
enc_aes_key = f.read(rsa_key.size_in_bytes())
nonce = f.read(16)
tag = f.read(16)
ciphertext = f.read()
aes_key = cipher_rsa.decrypt(enc_aes_key)
cipher_aes = AES.new(aes_key, AES.MODE_GCM, nonce=nonce)
data = cipher_aes.decrypt_and_verify(ciphertext, tag)
decrypted_name = self.decrypt_name(os.path.basename(encrypted_file_path), name_key, salt)
decrypted_file_path = os.path.join(os.path.dirname(encrypted_file_path), decrypted_name)
with open(decrypted_file_path, "wb") as file:
file.write(data)
os.remove(encrypted_file_path)
print(f"File {decrypted_file_path} has been decrypted.")
def encrypt_name(self, filename: str, key: str, salt: bytes) -> str:
"""
Encrypts the name of a file or directory.
:param salt: Salt for key derivation.
:param filename: The name to encrypt.
:param key: The key used for encryption.
"""
salted_key = self.derive_session_key(key, salt)
cipher = AES.new(salted_key, AES.MODE_GCM)
data = filename.encode("utf-8")
ciphertext, tag = cipher.encrypt_and_digest(data)
return base64.urlsafe_b64encode(cipher.nonce + tag + ciphertext).decode("utf-8")
def decrypt_name(self, encrypted_filename: str, key: str, salt: bytes) -> str:
"""
Decrypts the name of a file or directory.
:param salt: Salt for key derivation.
:param encrypted_filename: The encrypted name to decrypt.
:param key: The key used for decryption.
"""
salted_key = self.derive_session_key(key, salt)
data = base64.urlsafe_b64decode(encrypted_filename.encode("utf-8"))
nonce, tag, ciphertext = data[:16], data[16:32], data[32:]
cipher = AES.new(salted_key, AES.MODE_GCM, nonce=nonce)
return cipher.decrypt_and_verify(ciphertext, tag).decode("utf-8")
@staticmethod
def get_directory() -> str:
"""
Prompt users for directory name.
:return: Directory name
"""
while True:
directory = input("Enter the directory path to perform encryption/decryption: ").strip()
if not directory:
print("Error: Directory path cannot be empty.")
continue
if not os.path.exists(directory):
print(f"Error: The directory '{directory}' does not exist.")
continue
return directory
def run(self):
"""
Runs the program.
"""
while True:
try:
choice = int(input("Encrypt(0) or Decrypt(1)? "))
if choice not in [0, 1]:
raise ValueError(" Please enter 0 for encryption or 1 for decryption.")
break
except ValueError as e:
print(f"Invalid input: {e}")
if choice == 0:
self.generate_keys()
name_key = base64.b64encode(get_random_bytes(32)).decode("utf-8")
salt = get_random_bytes(16)
while True:
try:
choice = input("Do you want to secure your keys with a password? Enter only Y, y or N ,n: ").lower()
if choice not in ["y", "n"]:
raise ValueError("Please enter only Y, y or N, n.")
break
except ValueError as e:
print(f"Invalid input: {e}")
if choice == "n":
self.save_key(self.private_key, 'private_key.pem')
self.save_key(self.public_key, 'public_key.pem')
self.save_key(name_key.encode('utf-8'), 'name_key.pem', is_name_key=True, salt=salt)
elif choice == "y":
password = input("Enter your password here. Remember this password: ")
self.save_key(self.private_key, 'private_key.pem', password=password)
self.save_key(self.public_key, 'public_key.pem', password=password)
self.save_key(name_key.encode('utf-8'), 'name_key.pem', is_name_key=True, salt=salt)
directory = self.get_directory()
try:
encrypted_dir = self.encrypt_dir(directory, self.public_key, name_key, salt)
with open("dir.txt", "w") as file:
file.write(encrypted_dir)
print("Encryption completed.")
except Exception as e:
print(f"An exception occurred: {e}")
elif choice == 1:
while True:
try:
choice = input("Did you secure your keys with a password? Enter only Y, y or N ,n: ").lower()
if choice not in ["y", "n"]:
raise ValueError("Please enter only Y, y or N, n.")
break
except ValueError as e:
print(f"Invalid input: {e}")
if choice == "n":
loaded_private_key = self.load_key('private_key.pem')
salt, loaded_name_key = self.load_key('name_key.pem', is_name_key=True)
elif choice == "y":
password = input("Enter your password here: ")
try:
loaded_private_key = self.load_key('private_key.pem', password=password)
except ValueError as e:
print(f"Invalid password.")
sys.exit()
salt, loaded_name_key = self.load_key('name_key.pem', is_name_key=True)
with open("dir.txt", "r") as file:
dirs = file.read()
self.decrypt_dir(dirs, loaded_private_key, loaded_name_key, salt)
if __name__ == "__main__":
program = Program()
program.run()