Merge credsmanager branch #1
102
credsmanager.py
Normal file
102
credsmanager.py
Normal file
|
@ -0,0 +1,102 @@
|
||||||
|
from os.path import exists
|
||||||
|
from enum import Enum
|
||||||
|
from subprocess import Popen, PIPE
|
||||||
|
from cryptography.fernet import Fernet
|
||||||
|
import json
|
||||||
|
|
||||||
|
|
||||||
|
class ReturnCode(Enum):
|
||||||
|
UNKNOW_ERROR = 0
|
||||||
|
REG_KEY_ERROR = 1
|
||||||
|
ADD_REG_KEY_ERROR = 2
|
||||||
|
DEL_REG_KEY_ERROR = 3
|
||||||
|
INVALID_PATH = 4
|
||||||
|
REG_KEY_DELETED = 5
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
class Cryptography():
|
||||||
|
def __init__(self, appname: str):
|
||||||
|
self.appname = appname
|
||||||
|
|
||||||
|
key = self.__get_key()
|
||||||
|
if type(key) is not bytes:
|
||||||
|
key = self.__create_key()
|
||||||
|
if type(key) is not bytes:
|
||||||
|
raise ReturnCode.REG_KEY_ERROR
|
||||||
|
|
||||||
|
self.fernet = Fernet(key)
|
||||||
|
|
||||||
|
def __get_key(self):
|
||||||
|
try:
|
||||||
|
proces = Popen(f"reg QUERY HKEY_CURRENT_USER\{self.appname}", stdout=PIPE, stderr=PIPE) #Get key
|
||||||
|
if proces.stderr.read() !=b'':
|
||||||
|
return self.__create_key()
|
||||||
|
else:
|
||||||
|
return proces.stdout.read().decode("CP850").splitlines()[2].rsplit()[2].encode() #Return key as binary
|
||||||
|
|
||||||
|
except:
|
||||||
|
return ReturnCode.UNKNOW_ERROR
|
||||||
|
|
||||||
|
|
||||||
|
def __create_key(self):
|
||||||
|
try:
|
||||||
|
key = Fernet.generate_key()
|
||||||
|
proces = Popen(f"reg ADD HKEY_CURRENT_USER\{self.appname} /v fernkey /d {key.decode()}", stdout=PIPE, stderr=PIPE) #Add key in reg str
|
||||||
|
if proces.stderr.read() != b'':
|
||||||
|
return ReturnCode.ADD_REG_KEY_ERROR
|
||||||
|
return key
|
||||||
|
|
||||||
|
except:
|
||||||
|
return ReturnCode.REG_KEY_ERROR
|
||||||
|
|
||||||
|
|
||||||
|
def delete_key(self):
|
||||||
|
"""Delete the key used for decryption, WARNING ONCE THE KEY IS DELETED ALL YOUR CRYPTED DATA ARE UNDECRYPTABLE AND LOST"""
|
||||||
|
try:
|
||||||
|
proces = Popen(f"reg DELETE HKEY_CURRENT_USER\{self.appname} /f", stdout=PIPE, stderr=PIPE)
|
||||||
|
stderr = proces.stderr.read()
|
||||||
|
if stderr != b'':
|
||||||
|
return ReturnCode.DEL_REG_KEY_ERROR
|
||||||
|
return ReturnCode.REG_KEY_DELETED
|
||||||
|
|
||||||
|
except:
|
||||||
|
return ReturnCode.UNKNOW_ERROR
|
||||||
|
|
||||||
|
def decrypt(self, data: bytes):
|
||||||
|
return self.fernet.decrypt(data)
|
||||||
|
|
||||||
|
def encrypt(self, data: bytes):
|
||||||
|
return self.fernet.encrypt(data)
|
||||||
|
|
||||||
|
|
||||||
|
class Credentials():
|
||||||
|
def __init__(self, appname: str, credentials_path: str):
|
||||||
|
"""Used to decode crypted credential csv files"""
|
||||||
|
|
||||||
|
if not exists(self.creds_path):
|
||||||
|
return ReturnCode.INVALID_PATH
|
||||||
|
|
||||||
|
self.creds_path = credentials_path
|
||||||
|
self.crypt = Cryptography(appname)
|
||||||
|
|
||||||
|
|
||||||
|
def get_creds(self):
|
||||||
|
"""Return the uncrypted version of the usernames and passwords as a list of tuple : [(username0, password0), (username1, password1), ...]"""
|
||||||
|
with open(self.creds_path, "rb") as f:
|
||||||
|
return json.loads(self.crypt.decrypt(f.read()).decode())
|
||||||
|
|
||||||
|
|
||||||
|
def write_creds(self, data: list[tuple[str, str]]):
|
||||||
|
with open(self.creds_path, "wb") as f:
|
||||||
|
f.write(self.crypt.encrypt(json.dumps(data).encode()))
|
||||||
|
return
|
||||||
|
|
||||||
|
def add_cred(self, username: str, password: str):
|
||||||
|
"""Add username and password to the crypted file"""
|
||||||
|
|
||||||
|
data = self.get_creds()
|
||||||
|
data.append((username, password))
|
||||||
|
self.write_creds(data)
|
||||||
|
return
|
||||||
|
|
Loading…
Reference in a new issue