#!/usr/bin/env python3
"""Independent suite-1 composition oracle. Synthetic secrets ONLY; stdlib + system libsodium.
No Kotlin code or output is imported. Regenerate with: python3 tools/concealed-vectors/reference.py
"""
import ctypes as c
import ctypes.util
import hashlib
import json
import struct
from pathlib import Path
s = c.CDLL(ctypes.util.find_library('sodium'))
s.sodium_init()
def call(name, n, *args):
    out = c.create_string_buffer(n)
    assert getattr(s, name)(out, *args) == 0, name
    return out.raw

def h(x): return hashlib.blake2b(x,digest_size=32).digest()
def u32(x): return struct.pack('>I',x)
def u64(x): return struct.pack('>Q',x)
def v(x): return u32(len(x))+x
def label(x): return v(x.encode())
def kdf(n,k,l,ctx): return hashlib.blake2b(label('kiloparsec/concealed/1')+label(l)+v(ctx),key=k,digest_size=n).digest()
def box(m,n,k): return call('crypto_secretbox_easy',len(m)+16,m,c.c_ulonglong(len(m)),n,k)
def envelope(k,ctx,m):
    length=u32(len(m))
    return b''.join(box(part,kdf(24,k,'flight/'+name+'-iv',ctx),kdf(32,k,'flight/'+name+'-key',ctx)) for name,part in [('header',length),('body',length+m)])
def signkey(seed):
    pk=c.create_string_buffer(32); sk=c.create_string_buffer(64)
    assert s.crypto_sign_seed_keypair(pk,sk,seed)==0
    return pk.raw,sk.raw

def sign(m,sk):
    out=c.create_string_buffer(64); size=c.c_ulonglong()
    assert s.crypto_sign_detached(out,c.byref(size),m,c.c_ulonglong(len(m)),sk)==0
    return out.raw

def decode(r):
    p=2**255-19; A=486662; r=int.from_bytes(r,'little')&((1<<254)-1)
    w=-A*pow((1+2*r*r)%p,-1,p)%p
    square=pow((w**3+A*w*w+w)%p,(p-1)//2,p) in (0,1)
    return (w if square else (-w-A)%p).to_bytes(32,'little')

# Original ChaCha20 implementation: 64-bit counter and 64-bit nonce.
def chacha(key,nonce,data):
    def rot(x,n): return ((x<<n)|(x>>(32-n)))&0xffffffff
    def quarter(x,a,b,c_,d):
        x[a]=(x[a]+x[b])&0xffffffff; x[d]=rot(x[d]^x[a],16)
        x[c_]=(x[c_]+x[d])&0xffffffff; x[b]=rot(x[b]^x[c_],12)
        x[a]=(x[a]+x[b])&0xffffffff; x[d]=rot(x[d]^x[a],8)
        x[c_]=(x[c_]+x[d])&0xffffffff; x[b]=rot(x[b]^x[c_],7)
    out=bytearray()
    for counter in range((len(data)+63)//64):
        state=list(struct.unpack('<4I',b'expand 32-byte k'))+list(struct.unpack('<8I',key))+[counter&0xffffffff,counter>>32]+list(struct.unpack('<2I',nonce))
        x=state[:]
        for _ in range(10):
            for q in [(0,4,8,12),(1,5,9,13),(2,6,10,14),(3,7,11,15),(0,5,10,15),(1,6,11,12),(2,7,8,13),(3,4,9,14)]: quarter(x,*q)
        out.extend(struct.pack('<16I',*((a+b)&0xffffffff for a,b in zip(x,state))))
    return bytes(a^b for a,b in zip(data,out))

def main():
    f={}
    def put(name,value): f[name]=value.hex(); return value
    e=put('serverSecret',bytes(range(32)))
    E=put('serverPublic',call('crypto_scalarmult_base',32,e))
    cs=put('clientDhSecret',bytes(range(32,64))); C=put('clientDhPublic',call('crypto_scalarmult_base',32,cs))
    ss=put('serverDhSecret',bytes(range(64,96))); S=put('serverDhPublic',call('crypto_scalarmult_base',32,ss))
    vs,sks=signkey(put('serverSigningSeed',bytes(range(96,128)))); put('serverIdentity',vs)
    vc,skc=signkey(put('clientSigningSeed',bytes(range(128,160)))); put('clientIdentity',vc)
    R=put('representative',bytes(range(160,192))); U=put('heloIv',bytes(range(192,216)))
    shared=put('bootstrapShared',call('crypto_scalarmult',32,e,decode(R)))
    context=R+U+E; root=put('bootstrapRoot',kdf(32,shared,'bootstrap',context))
    bounds=u32(1048576)+struct.pack('>H',0)+u32(0)
    core=put('heloCore',bytes([1,1])+C+vs+bounds)
    F1=put('f1',R+U+envelope(root,h(context),v(core)))
    h1=put('h1',h(label('helo')+v(F1)))
    I=put('replyIv',bytes(range(24))); replyroot=put('replyRoot',kdf(32,C,'server-reply',h1+E+vs))
    core=put('serverCore',bytes([1,1])+S+h(label('helo-decrypted')+C)+vs+bounds)
    sig=put('serverSignature',sign(label('kiloparsec/concealed/1/server-auth')+h1+v(core),sks))
    F2=put('f2',I+envelope(replyroot,h(label('server-flight')+h1+I),v(core)+sig))
    h2=put('h2',h(label('handshake')+v(F1)+v(F2)))
    # Independent crypto_kx formula, not a crypto_kx API call.
    shared=put('sessionShared',call('crypto_scalarmult',32,cs,S))
    kx=hashlib.blake2b(shared+C+S,digest_size=64).digest()
    s2c=put('s2c',kx[:32]); c2s=put('c2s',kx[32:])
    token=put('token',kdf(32,c2s+s2c,'connection-token',h2))
    finishc=put('finishedC2s',kdf(32,c2s,'finished/c2s',h2)); finishs=put('finishedS2c',kdf(32,s2c,'finished/s2c',h2))
    sig=put('clientSignature',sign(label('kiloparsec/concealed/1/client-auth')+token+h2+vc,skc))
    core=put('ackCore',bytes([3])+h2+vc+sig); A=put('ackIv',bytes(range(24,48)))
    F3=put('f3',A+envelope(finishc,h(label('ack')+h2+A),v(core)))
    h3=put('h3',h(label('authenticated-handshake')+h2+v(F3)))
    core=put('acceptCore',bytes([4])+h3); J=put('acceptIv',bytes(range(48,72)))
    put('f4',J+envelope(finishs,h(label('accept')+h2+J),v(core)))
    for name,key in [('c2s',c2s),('s2c',s2c)]:
        for suffix,label_,n in [('StreamKey','stream',32),('StreamSalt','salt',24),('MacKey','mac',32),('UdpKey','udp',32)]: put(name+suffix,kdf(n,key,'traffic/'+label_+'/'+name,h3))
    get=lambda name: bytes.fromhex(f[name])
    streamkey=hashlib.blake2b(b'kiloparsec/concealed/stream-key/v1'+get('c2sStreamSalt'),key=get('c2sStreamKey'),digest_size=32).digest()
    streamnonce=hashlib.blake2b(b'kiloparsec/concealed/stream-nonce/v1'+get('c2sStreamSalt'),key=get('c2sStreamKey'),digest_size=16).digest()[:8]
    objects=[bytes([0,1,2]),bytes(range(70)),bytes([1,0,0])]; plain=b''
    for q,obj in enumerate(objects):
        put('object'+str(q),obj); body=v(obj); length=u32(len(body))
        tag=hashlib.blake2b(b'kiloparsec/concealed/record-mac/v1'+u64(q)+length+body,key=get('c2sMacKey'),digest_size=16).digest()
        plain+=length+body+tag
        iv=bytes((i+q*24)%256 for i in range(24))
        put('udp'+str(q),iv+box(u64(q)+body,iv,get('c2sUdpKey')))
    put('tcp',chacha(streamkey,streamnonce,plain))
    Path('tools/concealed-vectors/fixtures.json').write_text(json.dumps(f,indent=2)+'\n')
    kotlin='''@file:OptIn(ExperimentalUnsignedTypes::class)\npackage net.sergeych.kiloparsec.concealed.handshake\n\n// Generated by tools/concealed-vectors/reference.py; synthetic secrets ONLY.\ninternal object IndependentWireFixtures {\n    fun bytes(name: String): UByteArray = values.getValue(name).chunked(2).map { it.toUByte(16) }.toUByteArray()\n    private val values = mapOf(\n'''+''.join('        "'+k+'" to "'+value+'",\n' for k,value in f.items())+'    )\n}\n'
    Path('src/commonTest/kotlin/net/sergeych/kiloparsec/concealed/handshake/IndependentWireFixtures.kt').write_text(kotlin)
    print('Generated',len(f),'synthetic fixture fields')
if __name__=='__main__': main()