slakr/encryption.py at master · slakr-tech/slakr · GitHub

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
from cryptography.fernet import Fernet
import os, sys

encryption_key = os.environ.get('chatreEncryptionKey')

def encrypt(string, bts=True):
    if not os.environ.get('chatreEncryptionKey'):
        print('You must create an encryption key')

    if type(string) == str:
        encoded = string.encode()
    else:
        encoded = string

    f = Fernet(encryption_key)

    if bts:
        return f.encrypt(encoded)

    else:
        return f.encrypt(encoded).decode()


def decrypt(string, bts=False):
    if not os.environ.get('chatreEncryptionKey'):
        print('You must create an encryption key')

    encryption_key = os.environ.get('chatreEncryptionKey').encode()
    if type(string) == str:
        encoded = string.encode()
    else:
        encoded = string
    f = Fernet(encryption_key)
    decrypted = f.decrypt(encoded)

    if bts:
        return decrypted

    else:
        return decrypted.decode()


def test():
    string = 'test'
    if len(sys.argv) > 1:
        string = sys.argv[1]

    print(string)

    encrypted = encrypt(string)

    print(encrypted)

    print(decrypt(encrypted))

if __name__ == "__main__":
    test()