#!/usr/bin/env python3
import os,re,sys,base64,json,hashlib
from pathlib import Path
from cryptography import x509
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding,rsa,ec
TOKEN_RE=re.compile(rb'eyJ[A-Za-z0-9_-]{5,}\.[A-Za-z0-9_-]{2,}\.[A-Za-z0-9_-]*')
cert=x509.load_pem_x509_certificate(Path('work/scratch/layer7-audit/api-leaf.pem').read_bytes())
pub=cert.public_key()
seen={}; valid=[]; total=0
algmap={'RS256':hashes.SHA256(),'RS384':hashes.SHA384(),'RS512':hashes.SHA512(),'PS256':hashes.SHA256(),'PS384':hashes.SHA384(),'PS512':hashes.SHA512()}
for root,dirs,files in os.walk('work'):
    dirs[:] = [d for d in dirs if d not in {'.git','node_modules'}]
    for fn in files:
      p=os.path.join(root,fn)
      try:
        if os.path.getsize(p)>50_000_000: continue
        data=Path(p).read_bytes()
      except: continue
      for m in TOKEN_RE.finditer(data):
        tok=m.group().decode('ascii','ignore')
        if tok in seen: continue
        seen[tok]=p; total+=1
        try:
          h,pay,s=tok.split('.')
          hdr=json.loads(base64.urlsafe_b64decode(h+'='*(-len(h)%4)))
          alg=hdr.get('alg')
          if alg not in algmap or not s: continue
          sig=base64.urlsafe_b64decode(s+'='*(-len(s)%4))
          data2=(h+'.'+pay).encode()
          hv=algmap[alg]
          pad=padding.PKCS1v15() if alg.startswith('RS') else padding.PSS(mgf=padding.MGF1(hv),salt_length=hv.digest_size)
          pub.verify(sig,data2,pad,hv)
          payload=json.loads(base64.urlsafe_b64decode(pay+'='*(-len(pay)%4)))
          valid.append((p,alg,hdr,payload,tok))
        except Exception: pass
print('unique_tokens',total)
print('valid_with_api_tls_key',len(valid))
for p,alg,hdr,payload,tok in valid:
  print('\nFILE',p,'ALG',alg)
  print('HEADER',json.dumps(hdr,sort_keys=True))
  print('PAYLOAD',json.dumps(payload,sort_keys=True)[:2000])
  print('TOKEN',tok)