86 lines
2.6 KiB
Python
86 lines
2.6 KiB
Python
import hmac
|
|
from base64 import urlsafe_b64encode, urlsafe_b64decode
|
|
from hashlib import sha224, sha256, sha384, sha512
|
|
from ecdsa import ecdsa, SigningKey, VerifyingKey, NIST256p, NIST224p, NIST384p, NIST521p
|
|
|
|
allowed_curve = [
|
|
NIST224p,
|
|
NIST256p,
|
|
NIST384p,
|
|
NIST521p
|
|
]
|
|
|
|
hashfunc = [
|
|
sha224,
|
|
sha256,
|
|
sha384,
|
|
sha512
|
|
]
|
|
|
|
key_size = [28, 32, 48, 66]
|
|
signature_size = [56, 64, 96, 132]
|
|
|
|
|
|
class PastaSigner:
|
|
def __init__(self, secret: bytes, version: int):
|
|
self.purpose = 'public'
|
|
if version > 4 or version < 1:
|
|
version = 1
|
|
self.version = version
|
|
self.hashfunc = hashfunc[self.version - 1]
|
|
self.key_size = key_size[self.version - 1]
|
|
self.priv = SigningKey.from_string(secret[:self.key_size], curve=allowed_curve[self.version - 1])
|
|
|
|
def serialize(self, data: bytes, sig):
|
|
token = 'v' + str(self.version) + '.'
|
|
token += self.purpose + '.'
|
|
token += urlsafe_b64encode(data + sig).decode().replace('=', '')
|
|
return token
|
|
|
|
def sign(self, data: str):
|
|
data = data.encode()
|
|
pub = self.priv.get_verifying_key().to_string()
|
|
h = self.hashfunc(data + pub).digest()
|
|
nonce = hmac.new(self.priv.to_string(), data, self.hashfunc).hexdigest()
|
|
sig = self.priv.sign_digest(h, k=int(nonce, 16))
|
|
|
|
return self.serialize(data + pub, sig)
|
|
|
|
|
|
class PastaVerifier:
|
|
|
|
def __init__(self, secret):
|
|
self.purpose = 'public'
|
|
self.secret = secret
|
|
|
|
def deserialize(self, data: bytes):
|
|
try:
|
|
version, purpose, payload = data.split(b'.')
|
|
version = int(version.replace(b'v', b''))
|
|
if version > 4 or version < 1:
|
|
return False
|
|
self.version = version
|
|
self.hashfunc = hashfunc[self.version - 1]
|
|
self.key_size = key_size[self.version - 1]
|
|
self.priv = SigningKey.from_string(self.secret[:self.key_size], curve=allowed_curve[self.version - 1])
|
|
|
|
raw_data = urlsafe_b64decode(payload + (b'==' * 2))
|
|
size = signature_size[self.version - 1]
|
|
signature = raw_data[-size:]
|
|
public_key = raw_data[-size * 2:-size]
|
|
message = raw_data[:-size]
|
|
|
|
return message, public_key, signature
|
|
except Exception as e:
|
|
return False
|
|
|
|
def verify(self, token: str):
|
|
deserialized = self.deserialize(token.encode())
|
|
if deserialized:
|
|
message, _, signature = deserialized
|
|
h = self.hashfunc(message).digest()
|
|
verifier = self.priv.get_verifying_key()
|
|
return verifier.verify_digest(signature, h)
|
|
|
|
return False
|