-
Notifications
You must be signed in to change notification settings - Fork 1
/
set_api_key.py
74 lines (63 loc) · 2.26 KB
/
set_api_key.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
from cryptography.fernet import Fernet
import base64
import os
# File where the encrypted API key will be stored
ENCRYPTED_FILE = 'encrypted_api_key.bin'
def generate_key(pin):
"""Generate a key based on the pin code."""
pin_str = str(pin) # Convert pin to string
return base64.urlsafe_b64encode(pin_str.zfill(32).encode())
def setup_api_key(api_key, pin):
"""Encrypt and save the API key using the pin code."""
key = generate_key(pin)
fernet = Fernet(key)
encrypted_api_key = fernet.encrypt(api_key.encode())
with open(ENCRYPTED_FILE, 'wb') as file:
file.write(encrypted_api_key)
return "API key encrypted and saved."
def read_api_key(pin):
"""Read and decrypt the API key using the pin code."""
if not os.path.exists(ENCRYPTED_FILE):
return "No encrypted API key found."
key = generate_key(pin)
fernet = Fernet(key)
try:
with open(ENCRYPTED_FILE, 'rb') as file:
encrypted_api_key = file.read()
return fernet.decrypt(encrypted_api_key).decode()
except Exception:
return "Invalid pin or error in decryption."
def delete_api_key():
"""Delete the file containing the encrypted API key."""
if os.path.exists(ENCRYPTED_FILE):
os.remove(ENCRYPTED_FILE)
return "API key deleted."
else:
return "No API key to delete."
def main():
while True:
print("\nAPI Key Manager")
print("1. Set up a new API Key")
print("2. Read API Key")
print("3. Delete API Key")
print("4. Exit")
choice = input("Enter your choice (1-4): ")
if choice == '1':
api_key = input("Enter the API key: ")
pin = input("Enter a pin code for encryption: ")
response = setup_api_key(api_key, pin)
print(response)
elif choice == '2':
pin = input("Enter your pin code to decrypt the API key: ")
api_key = read_api_key(pin)
print("API Key:", api_key)
elif choice == '3':
response = delete_api_key()
print(response)
elif choice == '4':
print("Exiting.")
break
else:
print("Invalid choice. Please enter a number between 1 and 4.")
if __name__ == "__main__":
main()