scripts/tak1_reference.py
raw ยท 8665 bytes
"""Optional host-only reference runner for the supplied TAK decoder 2.3 DLL.
Install pefile and unicorn with pip. The Android app never uses this script,
the Windows DLL, or emulation. Only a small Win32 IO adapter is provided here;
this is not a general Windows runtime.
"""
import sys,struct
from pathlib import Path
import pefile
from unicorn import Uc,UC_ARCH_X86,UC_MODE_32,UC_HOOK_CODE,UC_HOOK_MEM_INVALID
from unicorn.x86_const import *
class Oracle:
def __init__(self,path):
self.pe=pefile.PE(path); self.base=self.pe.OPTIONAL_HEADER.ImageBase
self.uc=Uc(UC_ARCH_X86,UC_MODE_32); u=self.uc
u.mem_map(0,0x10000); u.mem_write(0,struct.pack('<I',0xffffffff))
u.mem_map(self.base,0x30000); u.mem_write(self.base,self.pe.get_memory_mapped_image())
u.mem_map(0x10000000,0x4000000); self.next=0x10001000
u.mem_map(0x60000000,0x100000); self.stack=0x600ff000
u.mem_map(0x70000000,0x10000); self.stop=0x70000000
self.funcs={};self.exports={s.name.decode():self.base+s.address for s in self.pe.DIRECTORY_ENTRY_EXPORT.symbols}
for imp in self.pe.DIRECTORY_ENTRY_IMPORT:
for ent in imp.imports:
a=0x70001000+len(self.funcs)*16;self.funcs[a]=ent.name.decode();u.mem_write(ent.address,struct.pack('<I',a));u.mem_write(a,b'\xc3')
self.files={};self.tls={};self.last=[];self.trace=False
u.hook_add(UC_HOOK_CODE,self.hook,begin=0x70000000,end=0x7000ffff)
u.hook_add(UC_HOOK_MEM_INVALID,self.invalid)
def read(self,a,n=4):return bytes(self.uc.mem_read(a,n))
def dword(self,a):return struct.unpack('<I',self.read(a))[0]
def write32(self,a,v):self.uc.mem_write(a,struct.pack('<I',v&0xffffffff))
def alloc(self,n):
a=self.next;self.next+=(max(n,1)+4095)&~4095
if self.next>=0x14000000:raise ValueError('heap exhausted')
return a
def cstr(self,a):
if not a:return ''
b=bytearray()
while (x:=self.read(a+len(b),1))!=b'\0':b+=x
return b.decode('latin1')
def ret(self,val,n):
sp=self.uc.reg_read(UC_X86_REG_ESP);ra=self.dword(sp)
self.uc.reg_write(UC_X86_REG_EAX,val&0xffffffff);self.uc.reg_write(UC_X86_REG_ESP,sp+4+n*4);self.uc.reg_write(UC_X86_REG_EIP,ra)
def hook(self,u,a,size,data):
self.last=(self.last+[a])[-15:]
if a==self.stop:u.emu_stop();return
if a not in self.funcs:return
name=self.funcs[a];sp=u.reg_read(UC_X86_REG_ESP);args=[self.dword(sp+4+i*4) for i in range(10)]
if self.trace:print(name,[hex(x) for x in args[:4]],flush=True)
if name in ('InitializeCriticalSection','EnterCriticalSection','LeaveCriticalSection','DeleteCriticalSection'):self.ret(1,1)
elif name=='VirtualAlloc':self.ret(self.alloc(args[1]),4)
elif name=='VirtualFree':self.ret(1,3)
elif name=='LocalAlloc':self.ret(self.alloc(args[1]),2)
elif name in ('LocalFree','CloseHandle','TlsFree','FreeLibrary'):self.ret(0 if name=='LocalFree' else 1,1)
elif name=='TlsAlloc':self.ret(1,0)
elif name=='TlsGetValue':self.ret(self.tls.get(args[0],0),1)
elif name=='TlsSetValue':self.tls[args[0]]=args[1];self.ret(1,2)
elif name=='GetVersionExA':u.mem_write(args[0]+4,struct.pack('<IIII',6,1,7601,2));self.ret(1,1)
elif name=='lstrcpynA':u.mem_write(args[0],self.cstr(args[1]).encode('latin1')[:max(0,args[2]-1)]+b'\0');self.ret(args[0],3)
elif name=='lstrlenA':self.ret(len(self.cstr(args[0])),1)
elif name=='LoadStringA':u.mem_write(args[2],b'Error\0');self.ret(5,4)
elif name=='LoadLibraryExA':self.ret(0,3)
elif name=='GetProcAddress':self.ret(0,2)
elif name=='GetSystemMetrics':self.ret(0,1)
elif name=='RegOpenKeyExA':self.ret(2,5)
elif name=='GetKeyboardType':self.ret(4 if args[0]==0 else 0,1)
elif name=='GetCurrentThreadId':self.ret(1,0)
elif name=='GetStartupInfoA':self.ret(0,1)
elif name in ('GetThreadLocale','GetACP'):self.ret(0x409 if name=='GetThreadLocale' else 1252,0)
elif name=='GetModuleHandleA':self.ret(self.base,1)
elif name=='GetCommandLineA':
p=self.alloc(100);u.mem_write(p,b'"C:\\tak.exe"\0');self.ret(p,0)
elif name=='GetModuleFileNameA':
s=b'C:\\tak_deco_lib.dll\0';u.mem_write(args[1],s);self.ret(len(s)-1,3)
elif name=='GetStdHandle':self.ret(0,1)
elif name=='GetCPInfo':u.mem_write(args[1],struct.pack('<I',1)+b'?\0'+b'\0'*14);self.ret(1,2)
elif name=='GetLocaleInfoA':
s=b'1252\0';u.mem_write(args[2],s[:args[3]]);self.ret(len(s),4)
elif name=='EnumCalendarInfoA':self.ret(1,4)
elif name=='GetLastError':self.ret(0,0)
elif name=='CreateFileA':
fn=self.cstr(args[0]);self.files[0x100]=open(self.input_path,'rb');self.ret(0x100,7)
elif name=='GetFileSize':
f=self.files[args[0]];p=f.tell();f.seek(0,2);n=f.tell();f.seek(p)
if args[1]:self.write32(args[1],n>>32)
self.ret(n,2)
elif name=='SetFilePointer':
f=self.files[args[0]];off=args[1] if args[1]<0x80000000 else args[1]-0x100000000
if args[2]:off+=(self.dword(args[2]) if self.dword(args[2])<0x80000000 else self.dword(args[2])-0x100000000)<<32
n=f.seek(off,args[3]);self.ret(n,4)
elif name=='ReadFile':
f=self.files[args[0]];b=f.read(args[2]);u.mem_write(args[1],b)
if args[3]:self.write32(args[3],len(b))
self.ret(1,5)
else:raise RuntimeError(f'Unhandled import {name}: {[hex(x) for x in args[:6]]}')
def invalid(self,u,access,a,size,value,data):
raise RuntimeError(f'DLL accessed unmapped memory at {a:#x}')
def call(self,fn,*args):
a=self.exports[fn] if isinstance(fn,str) else fn
u=self.uc;u.reg_write(UC_X86_REG_ESP,self.stack-4*(len(args)+1));u.mem_write(self.stack-4*(len(args)+1),struct.pack('<'+'I'*(len(args)+1),self.stop,*args))
u.reg_write(UC_X86_REG_EBP,0);u.reg_write(UC_X86_REG_EAX,0);u.reg_write(UC_X86_REG_EDX,0);u.reg_write(UC_X86_REG_ECX,0)
u.emu_start(a,self.stop,count=100000000)
if u.reg_read(UC_X86_REG_EIP)!=self.stop:raise RuntimeError('Instruction budget exceeded')
return u.reg_read(UC_X86_REG_EAX)
def open_reference(dll, source):
oracle = Oracle(str(dll))
oracle.call(oracle.base+oracle.pe.OPTIONAL_HEADER.AddressOfEntryPoint, oracle.base, 1, 0)
oracle.input_path = str(source)
filename = oracle.alloc(64)
oracle.uc.mem_write(filename, b"reference.tak\0")
options = oracle.alloc(64)
handle = oracle.call("tak_SSD_Create_FromFile", filename, options, 0, 0)
if not oracle.call("tak_SSD_Valid", handle):
raise RuntimeError(f"Reference DLL could not open the file; state {oracle.call('tak_SSD_State', handle):#x}")
return oracle, handle
def compare(dll, source, native, packets=None):
import subprocess, tempfile, hashlib
with tempfile.TemporaryDirectory(prefix="loopdeck-tak1-") as temp:
output = Path(temp)/"native.pcm"
command = [str(native), str(source), str(output)]
if packets is not None:
command.append(str(packets))
subprocess.run(command, check=True)
actual = output.read_bytes()
oracle, handle = open_reference(dll, source)
info = oracle.alloc(128)
if oracle.call("tak_SSD_GetStreamInfo", handle, info):
raise RuntimeError("Cannot read reference stream information")
block = oracle.dword(info+40)
if not block or len(actual)%block:
raise RuntimeError("Native output is not aligned to reference PCM frames")
destination, got = oracle.alloc(16384*block), oracle.alloc(4)
compared = 0
try:
while compared < len(actual):
frames = min(16384, (len(actual)-compared)//block)
status = oracle.call("tak_SSD_ReadAudio", handle, destination, frames, got)
if status or oracle.dword(got) != frames:
raise RuntimeError(f"Reference failed after {compared} bytes; status {status:#x}")
reference = oracle.read(destination, frames*block)
chunk = actual[compared:compared+len(reference)]
if chunk != reference:
offset = next(i for i,(a,b) in enumerate(zip(chunk,reference)) if a != b)
raise RuntimeError(f"PCM mismatch at byte {compared+offset}")
compared += len(reference)
finally:
oracle.call("tak_SSD_Destroy", handle)
for file in oracle.files.values():
file.close()
print(f"Exact PCM match: {compared} bytes; sha256 {hashlib.sha256(actual).hexdigest()}")
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--dll", type=Path, required=True)
parser.add_argument("--input", type=Path, required=True)
parser.add_argument("--native", type=Path, required=True)
parser.add_argument("--packets", type=int, help="Limit both sides to complete packets in a truncated sample")
args = parser.parse_args()
compare(args.dll.resolve(), args.input.resolve(), args.native.resolve(), args.packets)