-
Notifications
You must be signed in to change notification settings - Fork 0
/
aes256.py
65 lines (46 loc) · 1.84 KB
/
aes256.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
#this is an algorithm for 256 bit AES with custom padding and initialisation vector function
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
from Crypto.Util.Padding import unpad
from Crypto import Random
from Crypto.Hash import SHA256
import base64
import sys
import argparse
# padding used -> PKCS7
def encrypt(key, message, keytype):
message = message.encode()
if keytype == "hex":
key = bytes(bytearray.fromhex(key))
else:
key = key.encode()
key = SHA256.new(key).digest()
init_vector = Random.new().read(AES.block_size)
cipher = AES.new(key,AES.MODE_CBC, init_vector)
data = init_vector + cipher.encrypt(pad(message,AES.block_size)) # maintaining access to IV
return base64.b64encode(data).decode()
# return data
def decrypt(key, message, keytype):
message = base64.b64decode(message)
if keytype == "hex":
key = bytes(bytearray.fromhex(key))
else:
key = key.encode()
key = SHA256.new(key).digest()
init_vector = message[:AES.block_size] #finally the use of slicing in python
decryptor = AES.new(key, AES.MODE_CBC, init_vector)
data = unpad(decryptor.decrypt(message),AES.block_size)
final = data[AES.block_size:]
return final
if __name__ == "__main__":
parser= argparse.ArgumentParser(description = 'command for encryption/decryption');
parser.add_argument('-e','--encrypt', action='store_true')
parser.add_argument('-d','--decrypt', action='store_true')
parser.add_argument('-k', '--key', type=str)
parser.add_argument('-ky', '--keytype', type=str)
parser.add_argument('-m', '--message', type=str)
args=parser.parse_args()
if args.encrypt:
print(encrypt(args.key, args.message, args.keytype))
elif args.decrypt:
print(decrypt(args.key, args.message, args.keytype))