1
+ import pathlib
2
+ import secrets
3
+ import os
4
+ import base64
5
+ import getpass
6
+
7
+ import cryptography
8
+ from cryptography .fernet import Fernet
9
+ from cryptography .hazmat .primitives .kdf .scrypt import Scrypt
10
+
11
+
12
+ def generate_salt (size = 16 ):
13
+ """Generate the salt used for key derivation,
14
+ `size` is the length of the salt to generate"""
15
+ return secrets .token_bytes (size )
16
+
17
+
18
+ def derive_key (salt , password ):
19
+ """Derive the key from the `password` using the passed `salt`"""
20
+ kdf = Scrypt (salt = salt , length = 32 , n = 2 ** 14 , r = 8 , p = 1 )
21
+ return kdf .derive (password .encode ())
22
+
23
+
24
+ def load_salt ():
25
+ # load salt from salt.salt file
26
+ return open ("salt.salt" , "rb" ).read ()
27
+
28
+
29
+ def generate_key (password , salt_size = 16 , load_existing_salt = False , save_salt = True ):
30
+ """Generates a key from a `password` and the salt.
31
+ If `load_existing_salt` is True, it'll load the salt from a file
32
+ in the current directory called "salt.salt".
33
+ If `save_salt` is True, then it will generate a new salt
34
+ and save it to "salt.salt" """
35
+ if load_existing_salt :
36
+ # load existing salt
37
+ salt = load_salt ()
38
+ elif save_salt :
39
+ # generate new salt and save it
40
+ salt = generate_salt (salt_size )
41
+ with open ("salt.salt" , "wb" ) as salt_file :
42
+ salt_file .write (salt )
43
+ # generate the key from the salt and the password
44
+ derived_key = derive_key (salt , password )
45
+ # encode it using Base 64 and return it
46
+ return base64 .urlsafe_b64encode (derived_key )
47
+
48
+
49
+ def encrypt (filename , key ):
50
+ """Given a filename (str) and key (bytes), it encrypts the file and write it"""
51
+ f = Fernet (key )
52
+ with open (filename , "rb" ) as file :
53
+ # read all file data
54
+ file_data = file .read ()
55
+ # encrypt data
56
+ encrypted_data = f .encrypt (file_data )
57
+ # write the encrypted file
58
+ with open (filename , "wb" ) as file :
59
+ file .write (encrypted_data )
60
+
61
+
62
+ def encrypt_folder (foldername , key ):
63
+ # if it's a folder, encrypt the entire folder (i.e all the containing files)
64
+ for child in pathlib .Path (foldername ).glob ("*" ):
65
+ if child .is_file ():
66
+ print (f"[*] Encrypting { child } " )
67
+ # encrypt the file
68
+ encrypt (child , key )
69
+ elif child .is_dir ():
70
+ # if it's a folder, encrypt the entire folder by calling this function recursively
71
+ encrypt_folder (child , key )
72
+
73
+
74
+ def decrypt (filename , key ):
75
+ """Given a filename (str) and key (bytes), it decrypts the file and write it"""
76
+ f = Fernet (key )
77
+ with open (filename , "rb" ) as file :
78
+ # read the encrypted data
79
+ encrypted_data = file .read ()
80
+ # decrypt data
81
+ try :
82
+ decrypted_data = f .decrypt (encrypted_data )
83
+ except cryptography .fernet .InvalidToken :
84
+ print ("[!] Invalid token, most likely the password is incorrect" )
85
+ return
86
+ # write the original file
87
+ with open (filename , "wb" ) as file :
88
+ file .write (decrypted_data )
89
+
90
+
91
+ def decrypt_folder (foldername , key ):
92
+ # if it's a folder, decrypt the entire folder
93
+ for child in pathlib .Path (foldername ).glob ("*" ):
94
+ if child .is_file ():
95
+ print (f"[*] Decrypting { child } " )
96
+ # decrypt the file
97
+ decrypt (child , key )
98
+ elif child .is_dir ():
99
+ # if it's a folder, decrypt the entire folder by calling this function recursively
100
+ decrypt_folder (child , key )
101
+
102
+
103
+ if __name__ == "__main__" :
104
+ import argparse
105
+ parser = argparse .ArgumentParser (description = "File Encryptor Script with a Password" )
106
+ parser .add_argument ("path" , help = "Path to encrypt/decrypt, can be a file or an entire folder" )
107
+ parser .add_argument ("-s" , "--salt-size" , help = "If this is set, a new salt with the passed size is generated" ,
108
+ type = int )
109
+ parser .add_argument ("-e" , "--encrypt" , action = "store_true" ,
110
+ help = "Whether to encrypt the file/folder, only -e or -d can be specified." )
111
+ parser .add_argument ("-d" , "--decrypt" , action = "store_true" ,
112
+ help = "Whether to decrypt the file/folder, only -e or -d can be specified." )
113
+ # parse the arguments
114
+ args = parser .parse_args ()
115
+ # get the password
116
+ if args .encrypt :
117
+ password = getpass .getpass ("Enter the password for encryption: " )
118
+ elif args .decrypt :
119
+ password = getpass .getpass ("Enter the password you used for encryption: " )
120
+ # generate the key
121
+ if args .salt_size :
122
+ key = generate_key (password , salt_size = args .salt_size , save_salt = True )
123
+ else :
124
+ key = generate_key (password , load_existing_salt = True )
125
+ # get the encrypt and decrypt flags
126
+ encrypt_ = args .encrypt
127
+ decrypt_ = args .decrypt
128
+ # check if both encrypt and decrypt are specified
129
+ if encrypt_ and decrypt_ :
130
+ raise TypeError ("Please specify whether you want to encrypt the file or decrypt it." )
131
+ elif encrypt_ :
132
+ if os .path .isfile (args .path ):
133
+ # if it is a file, encrypt it
134
+ encrypt (args .path , key )
135
+ elif os .path .isdir (args .path ):
136
+ encrypt_folder (args .path , key )
137
+ elif decrypt_ :
138
+ if os .path .isfile (args .path ):
139
+ decrypt (args .path , key )
140
+ elif os .path .isdir (args .path ):
141
+ decrypt_folder (args .path , key )
142
+ else :
143
+ raise TypeError ("Please specify whether you want to encrypt the file or decrypt it." )
0 commit comments