import sys,re
sys.path.insert(0,'.')
from zb import *
TCLK=b'ZigBeeAlliance09'
keys={'tclk':TCLK,'kt':hmac_mmo(TCLK,b'\x00'),'kl':hmac_mmo(TCLK,b'\x02')}
nwkkeys=[]
def try_aux(buf,off,keyset,aadstart):
    sc=buf[off]; fc=int.from_bytes(buf[off+1:off+5],'little'); p=off+5
    src=None
    if sc&0x20: src=buf[p:p+8]; p+=8
    kid=(sc>>3)&3
    ks=None
    if kid==1: ks=buf[p]; p+=1
    sc5=(sc&~7)|5
    aad=bytes(buf[aadstart:off])+bytes([sc5])+bytes(buf[off+1:p])
    return sc5,kid,fc,src,ks,p,aad
for line in open(sys.argv[1]):
    t,st,*r=line.split()
    if st!='OK' or len(r)<2: continue
    fr=bytes.fromhex(r[1])[:-2]
    fcf=int.from_bytes(fr[:2],'little')
    if fcf&7!=1: continue
    # MAC hdr: panid compression assumed
    dm=(fcf>>10)&3; sm=(fcf>>14)&3
    p=3+2+(2 if dm==2 else 8)+(2 if sm==2 else 8)
    nwk=fr[p:]
    nf=int.from_bytes(nwk[:2],'little')
    q=8
    if nf&0x0800: q+=8
    if nf&0x1000: q+=8
    if nf&0x0400: print('multicast?')
    if nf&0x0200:
        sc,kid,fc,src,ks,pp,aad=try_aux(nwk,q,None,0)
        enc=nwk[pp:]
        ok=None
        for k in nwkkeys:
            try: ok=ccm_dec(k,src,fc,sc,aad,enc); break
            except Exception: pass
        if ok is None: print(t,'NWK enc, no key', src.hex(), fc); continue
        aps=ok
    else:
        aps=nwk[q:]
    print(t,'NWK src',nwk[4:6][::-1].hex(),'dst',nwk[2:4][::-1].hex(),'APS',aps.hex())
    afc=aps[0]
    if afc&0x20:
        typ=afc&3
        a=1
        if typ==1: pass  # command
        elif typ==0:
            a=1
            if (afc>>2)&3!=2: a+=1  # dst endpoint (unicast)
            a+=2+2+1  # cluster, profile, src ep
        a+=1  # counter
        sc,kid,fc,src,ks,pp,aad=try_aux(aps,a,None,0)
        enc=aps[pp:]
        print('   APS sec kid',kid,'fc',fc,'src',src.hex() if src else None)
        for name,k in keys.items():
            try:
                m=ccm_dec(k,src if src else bytes(8),fc,sc,aad,enc)
                print('   decrypted with',name,m.hex())
                if m[0]==5 and m[1]==1:
                    nk=m[2:18]; print('   NETWORK KEY',nk.hex()); nwkkeys.append(nk)
                print('   ascii',m)
            except Exception as e: pass