You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
75 lines
2.2 KiB
75 lines
2.2 KiB
4 years ago
|
import os
|
||
|
|
import re
|
||
|
|
import base64
|
||
|
|
from Crypto.PublicKey import RSA
|
||
|
|
from Crypto.Util import number
|
||
|
|
from typing import Any
|
||
|
|
from typing import Dict
|
||
|
|
from typing import Optional
|
||
|
|
|
||
|
|
KEY_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), '.')
|
||
|
|
|
||
|
|
class Key(object):
|
||
|
|
DEFAULT_KEY_SIZE = 2048
|
||
|
|
|
||
|
|
def __init__(self, owner: str) -> None:
|
||
|
|
self.owner = owner
|
||
|
|
self.privkey_pem: Optional[str] = None
|
||
|
|
self.pubkey_pem: Optional[str] = None
|
||
|
|
self.privkey: Optional[Any] = None
|
||
|
|
self.pubkey: Optional[Any] = None
|
||
|
|
|
||
|
|
def load_pub(self, pubkey_pem: str) -> None:
|
||
|
|
self.pubkey_pem = pubkey_pem
|
||
|
|
self.pubkey = RSA.importKey(pubkey_pem)
|
||
|
|
|
||
|
|
def load(self, privkey_pem: str) -> None:
|
||
|
|
self.privkey_pem = privkey_pem
|
||
|
|
self.privkey = RSA.importKey(self.privkey_pem)
|
||
|
|
self.pubkey_pem = self.privkey.publickey().exportKey("PEM").decode("utf-8")
|
||
|
|
|
||
|
|
def new(self) -> None:
|
||
|
|
k = RSA.generate(self.DEFAULT_KEY_SIZE)
|
||
|
|
self.privkey_pem = k.exportKey("PEM").decode("utf-8")
|
||
|
|
self.pubkey_pem = k.publickey().exportKey("PEM").decode("utf-8")
|
||
|
|
self.privkey = k
|
||
|
|
|
||
|
|
def key_id(self) -> str:
|
||
|
|
return f"{self.owner}#main-key"
|
||
|
|
|
||
|
|
def to_dict(self) -> Dict[str, Any]:
|
||
|
|
return {
|
||
|
|
"id": self.key_id(),
|
||
|
|
"owner": self.owner,
|
||
|
|
"publicKeyPem": self.pubkey_pem,
|
||
|
|
}
|
||
|
|
|
||
|
|
def to_magic_key(self) -> str:
|
||
|
|
mod = base64.urlsafe_b64encode(
|
||
|
|
number.long_to_bytes(self.privkey.n) # type: ignore
|
||
|
|
).decode("utf-8")
|
||
|
|
pubexp = base64.urlsafe_b64encode(
|
||
|
|
number.long_to_bytes(self.privkey.e) # type: ignore
|
||
|
|
).decode("utf-8")
|
||
|
|
return f"data:application/magic-public-key,RSA.{mod}.{pubexp}"
|
||
|
|
|
||
|
|
|
||
|
|
def get_key(owner: str) -> Key:
|
||
|
|
""""Loads or generates an RSA key."""
|
||
|
|
k = Key(owner)
|
||
|
|
user = re.sub('[^\w\d]', "_", owner)
|
||
|
|
key_path = os.path.join(KEY_DIR, f"key_{user}.pem")
|
||
|
|
if os.path.isfile(key_path):
|
||
|
|
with open(key_path) as f:
|
||
|
|
privkey_pem = f.read()
|
||
|
|
k.load(privkey_pem)
|
||
|
|
else:
|
||
|
|
k.new()
|
||
|
|
with open(key_path, "w") as f:
|
||
|
|
f.write(k.privkey_pem)
|
||
|
|
|
||
|
|
return k
|
||
|
|
|
||
|
|
|
||
|
|
|