137 lines
4.1 KiB
Python
137 lines
4.1 KiB
Python
import base64
|
|
import ctypes
|
|
import json
|
|
import os
|
|
import time
|
|
from ctypes import wintypes
|
|
|
|
|
|
class _DataBlob(ctypes.Structure):
|
|
_fields_ = [("cbData", wintypes.DWORD), ("pbData", ctypes.POINTER(ctypes.c_byte))]
|
|
|
|
|
|
_crypt32 = None
|
|
_kernel32 = None
|
|
if os.name == "nt":
|
|
_crypt32 = ctypes.WinDLL("crypt32", use_last_error=True)
|
|
_kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
|
|
_crypt32.CryptProtectData.argtypes = [
|
|
ctypes.POINTER(_DataBlob),
|
|
wintypes.LPCWSTR,
|
|
ctypes.POINTER(_DataBlob),
|
|
ctypes.c_void_p,
|
|
ctypes.c_void_p,
|
|
wintypes.DWORD,
|
|
ctypes.POINTER(_DataBlob),
|
|
]
|
|
_crypt32.CryptProtectData.restype = wintypes.BOOL
|
|
_crypt32.CryptUnprotectData.argtypes = [
|
|
ctypes.POINTER(_DataBlob),
|
|
ctypes.POINTER(wintypes.LPWSTR),
|
|
ctypes.POINTER(_DataBlob),
|
|
ctypes.c_void_p,
|
|
ctypes.c_void_p,
|
|
wintypes.DWORD,
|
|
ctypes.POINTER(_DataBlob),
|
|
]
|
|
_crypt32.CryptUnprotectData.restype = wintypes.BOOL
|
|
|
|
|
|
def _crypt_protect_data(data, description=""):
|
|
buffer = ctypes.create_string_buffer(data)
|
|
data_in = _DataBlob(len(data), ctypes.cast(buffer, ctypes.POINTER(ctypes.c_byte)))
|
|
data_out = _DataBlob()
|
|
if not _crypt32.CryptProtectData(ctypes.byref(data_in), description, None, None, None, 0, ctypes.byref(data_out)):
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
try:
|
|
return ctypes.string_at(data_out.pbData, data_out.cbData)
|
|
finally:
|
|
_kernel32.LocalFree(data_out.pbData)
|
|
|
|
|
|
def _crypt_unprotect_data(data):
|
|
buffer = ctypes.create_string_buffer(data)
|
|
data_in = _DataBlob(len(data), ctypes.cast(buffer, ctypes.POINTER(ctypes.c_byte)))
|
|
data_out = _DataBlob()
|
|
if not _crypt32.CryptUnprotectData(ctypes.byref(data_in), None, None, None, None, 0, ctypes.byref(data_out)):
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
try:
|
|
return ctypes.string_at(data_out.pbData, data_out.cbData)
|
|
finally:
|
|
_kernel32.LocalFree(data_out.pbData)
|
|
|
|
|
|
def _encrypt_token(token):
|
|
if os.name != "nt":
|
|
raise RuntimeError("DPAPI is available only on Windows.")
|
|
encrypted_bytes = _crypt_protect_data(token.encode("utf-8"))
|
|
return base64.b64encode(encrypted_bytes).decode("ascii")
|
|
|
|
|
|
def _decrypt_token(token_data):
|
|
if os.name != "nt":
|
|
raise RuntimeError("DPAPI is available only on Windows.")
|
|
encrypted_bytes = base64.b64decode(token_data.encode("ascii"))
|
|
decrypted_bytes = _crypt_unprotect_data(encrypted_bytes)
|
|
return decrypted_bytes.decode("utf-8")
|
|
|
|
|
|
def save_token(token, token_file, app_data_dir, expires_in=0):
|
|
try:
|
|
expires_in = int(expires_in)
|
|
except (ValueError, TypeError):
|
|
expires_in = 0
|
|
|
|
os.makedirs(app_data_dir, exist_ok=True)
|
|
expiration_time = (time.time() + expires_in) if expires_in > 0 else 0
|
|
|
|
stored_token = token
|
|
encrypted = False
|
|
if os.name == "nt":
|
|
try:
|
|
stored_token = _encrypt_token(token)
|
|
encrypted = True
|
|
except Exception as exc:
|
|
raise RuntimeError("Failed to securely store token with DPAPI.") from exc
|
|
|
|
data = {
|
|
"token": stored_token,
|
|
"expiration_time": expiration_time,
|
|
"encrypted": encrypted,
|
|
}
|
|
|
|
with open(token_file, "w", encoding="utf-8") as f:
|
|
json.dump(data, f)
|
|
return expiration_time
|
|
|
|
|
|
def load_token(token_file):
|
|
if not os.path.exists(token_file):
|
|
return None, None
|
|
|
|
with open(token_file, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
|
|
token = data.get("token")
|
|
encrypted = data.get("encrypted", False)
|
|
if token and encrypted:
|
|
try:
|
|
token = _decrypt_token(token)
|
|
except Exception:
|
|
try:
|
|
os.remove(token_file)
|
|
except Exception:
|
|
pass
|
|
return None, None
|
|
|
|
expiration_time = data.get("expiration_time")
|
|
if token and (expiration_time == 0 or expiration_time > time.time()):
|
|
return token, expiration_time
|
|
|
|
try:
|
|
os.remove(token_file)
|
|
except Exception:
|
|
pass
|
|
return None, None
|
|
|