tai/taicpp/tests/whatsapp_data/appstate_oracle.go

// Oracle generator for detz::wa app-state sync vectors.
//
// Every formula is copied from whatsmeow (MPL-2.0) appstate/keys.go, hash.go,
// decode.go and encode of the lthash package, and every wire message is
// produced by whatsmeow's own protobuf types, so the C++ decoder is checked
// against the reference bytes, not against our own encoder.
//
// Run from a module that requires go.mau.fi/whatsmeow:
//   go run appstate_oracle.go > appstate_vectors.json (inside a module that requires go.mau.fi/whatsmeow)
//go:build ignore

package main

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"hash"
	"os"

	"go.mau.fi/whatsmeow/appstate/lthash"
	"go.mau.fi/whatsmeow/proto/waE2E"
	"go.mau.fi/whatsmeow/proto/waServerSync"
	"go.mau.fi/whatsmeow/proto/waSyncAction"
	"go.mau.fi/whatsmeow/util/hkdfutil"
	"google.golang.org/protobuf/proto"
)

const collection = "critical_unblock_low"

type keys struct{ index, valueEnc, valueMAC, snapshotMAC, patchMAC []byte }

func expand(keyData []byte) keys {
	e := hkdfutil.SHA256(keyData, nil, []byte("WhatsApp Mutation Keys"), 160)
	return keys{e[0:32], e[32:64], e[64:96], e[96:128], e[128:160]}
}

func u64(v uint64) []byte { b := make([]byte, 8); binary.BigEndian.PutUint64(b, v); return b }

func concatHMAC(alg func() hash.Hash, key []byte, data ...[]byte) []byte {
	h := hmac.New(alg, key)
	for _, d := range data {
		h.Write(d)
	}
	return h.Sum(nil)
}

func contentMAC(op int, data, keyID, key []byte) []byte {
	return concatHMAC(sha512.New, key, []byte{byte(op) + 1}, keyID, data, u64(uint64(len(keyID)+1)))[:32]
}

func aesCBCEncrypt(key, iv, pt []byte) []byte {
	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}
	pad := 16 - len(pt)%16
	padded := append(append([]byte{}, pt...), bytes.Repeat([]byte{byte(pad)}, pad)...)
	out := make([]byte, len(padded))
	cipher.NewCBCEncrypter(block, iv).CryptBlocks(out, padded)
	return out
}

type entry struct {
	op      int // 0 SET, 1 REMOVE
	index   []string
	action  *waSyncAction.SyncActionValue
	valMAC  []byte // value MAC of this entry (SET only, used for the hash)
	idxMAC  []byte
	valBlob []byte
}

var ivSeed byte = 1

func encrypt(k keys, keyID []byte, op int, index []string, action *waSyncAction.SyncActionValue) entry {
	idxJSON, _ := json.Marshal(index)
	sad := &waSyncAction.SyncActionData{Index: idxJSON, Value: action, Version: proto.Int32(2)}
	pt, err := proto.Marshal(sad)
	if err != nil {
		panic(err)
	}
	iv := bytes.Repeat([]byte{ivSeed}, 16)
	ivSeed++
	ct := aesCBCEncrypt(k.valueEnc, iv, pt)
	content := append(append([]byte{}, iv...), ct...)
	vmac := contentMAC(op, content, keyID, k.valueMAC)
	blob := append(append([]byte{}, content...), vmac...)
	idxMAC := concatHMAC(sha256.New, k.index, idxJSON)
	return entry{op: op, index: index, action: action, valMAC: vmac, idxMAC: idxMAC, valBlob: blob}
}

func rec(e entry, keyID []byte) *waServerSync.SyncdRecord {
	return &waServerSync.SyncdRecord{
		Index:  &waServerSync.SyncdIndex{Blob: e.idxMAC},
		Value:  &waServerSync.SyncdValue{Blob: e.valBlob},
		KeyID:  &waServerSync.KeyId{ID: keyID},
	}
}

func contact(full, first string) *waSyncAction.SyncActionValue {
	return &waSyncAction.SyncActionValue{
		Timestamp:     proto.Int64(1767225600000),
		ContactAction: &waSyncAction.ContactAction{FullName: proto.String(full), FirstName: proto.String(first)},
	}
}

type hashState struct {
	version uint64
	hash    [128]byte
}

func (hs *hashState) apply(removed, added [][]byte) {
	lthash.WAPatchIntegrity.SubtractThenAddInPlace(hs.hash[:], removed, added)
}

func valueMACOf(e entry) []byte { return e.valMAC }

func main() {
	keyID := []byte{0x00, 0x00, 0x00, 0x01, 0x42, 0x10}
	keyData := make([]byte, 32)
	for i := range keyData {
		keyData[i] = byte(0x11 + i)
	}
	k := expand(keyData)

	// ---- snapshot v1: four records, all SET ----
	A := encrypt(k, keyID, 0, []string{"contact", "31611111111@s.whatsapp.net"}, contact("Nigel Dokter", "Nigel"))
	B := encrypt(k, keyID, 0, []string{"contact", "31622222222@s.whatsapp.net"}, contact("", "Ann"))
	C := encrypt(k, keyID, 0, []string{"contact", "31633333333@s.whatsapp.net"}, contact("Pieter Bos", "Pieter"))
	D := encrypt(k, keyID, 0, []string{"setting_pushName"}, &waSyncAction.SyncActionValue{Timestamp: proto.Int64(1767225600000), PushNameSetting: &waSyncAction.PushNameSetting{Name: proto.String("detz")}})
	snapEntries := []entry{A, B, C, D}

	state := hashState{version: 1}
	var added [][]byte
	var snapRecs []*waServerSync.SyncdRecord
	for _, e := range snapEntries {
		added = append(added, e.valMAC)
		snapRecs = append(snapRecs, rec(e, keyID))
	}
	state.apply(nil, added)
	snapMAC := concatHMAC(sha256.New, k.snapshotMAC, state.hash[:], u64(state.version), []byte(collection))
	snapshot := &waServerSync.SyncdSnapshot{
		Version: &waServerSync.SyncdVersion{Version: proto.Uint64(state.version)},
		Records: snapRecs,
		Mac:     snapMAC,
		KeyID:   &waServerSync.KeyId{ID: keyID},
	}
	snapBytes, err := proto.Marshal(snapshot)
	if err != nil {
		panic(err)
	}
	snapHash := hex.EncodeToString(state.hash[:])

	// previous value MACs by index MAC, as the store would hold them after the snapshot
	prev := map[string][]byte{}
	for _, e := range snapEntries {
		prev[hex.EncodeToString(e.idxMAC)] = e.valMAC
	}

	// ---- patch v2 ----
	// m1: SET A again (new name): removes A's snapshot value MAC, adds the new one
	A2 := encrypt(k, keyID, 0, []string{"contact", "31611111111@s.whatsapp.net"}, contact("Nigel Dokter (work)", "Nigel"))
	// m2: REMOVE C: removes C's snapshot value MAC
	C2 := encrypt(k, keyID, 1, []string{"contact", "31633333333@s.whatsapp.net"}, &waSyncAction.SyncActionValue{Timestamp: proto.Int64(1767225600000)})
	// m3: SET E with first name only
	E1 := encrypt(k, keyID, 0, []string{"contact", "31644444444@s.whatsapp.net"}, contact("", "Joop"))
	// m4: SET E again (in-patch previous value, no DB lookup)
	E2 := encrypt(k, keyID, 0, []string{"contact", "31644444444@s.whatsapp.net"}, contact("Joop Visser", "Joop"))
	// m5: SET F with neither name (must be stored as unnamed = skipped by the contact sync)
	F := encrypt(k, keyID, 0, []string{"contact", "31655555555@s.whatsapp.net"}, contact("", ""))
	// m6: a non-contact SET with an index the contact sync must ignore
	M := encrypt(k, keyID, 0, []string{"setting_pushName"}, &waSyncAction.SyncActionValue{Timestamp: proto.Int64(1767225600001), PushNameSetting: &waSyncAction.PushNameSetting{Name: proto.String("detz2")}})

	patchEntries := []entry{A2, C2, E1, E2, F, M}
	var removed [][]byte
	var padded [][]byte
	_ = padded
	inPatch := map[string]int{} // idxMAC -> index of latest mutation in this patch that touched it
	var addedP [][]byte
	for i, e := range patchEntries {
		key := hex.EncodeToString(e.idxMAC)
		if j, ok := inPatch[key]; ok {
			if patchEntries[j].op == 0 {
				removed = append(removed, patchEntries[j].valMAC)
			}
		} else if pv, ok := prev[key]; ok {
			removed = append(removed, pv)
		}
		inPatch[key] = i
		if e.op == 0 {
			addedP = append(addedP, e.valMAC)
		}
	}
	// the prior value of an index is consumed once; mirror whatsmeow: removal = prev SET value
	state.version = 2
	state.apply(removed, addedP)

	patchSnapMAC := concatHMAC(sha256.New, k.snapshotMAC, state.hash[:], u64(state.version), []byte(collection))
	var muts []*waServerSync.SyncdMutation
	var valMACs [][]byte
	for _, e := range patchEntries {
		op := waServerSync.SyncdMutation_SET
		if e.op == 1 {
			op = waServerSync.SyncdMutation_REMOVE
		}
		muts = append(muts, &waServerSync.SyncdMutation{Operation: op.Enum(), Record: rec(e, keyID)})
		valMACs = append(valMACs, e.valMAC)
	}
	// patch MAC: snapshotMAC || value MACs || version || name
	pmData := [][]byte{patchSnapMAC}
	pmData = append(pmData, valMACs...)
	pmData = append(pmData, u64(state.version), []byte(collection))
	patchMAC := concatHMAC(sha256.New, k.patchMAC, pmData...)
	patch := &waServerSync.SyncdPatch{
		Version:     &waServerSync.SyncdVersion{Version: proto.Uint64(state.version)},
		Mutations:   muts,
		SnapshotMAC: patchSnapMAC,
		PatchMAC:    patchMAC,
		KeyID:       &waServerSync.KeyId{ID: keyID},
	}
	patchBytes, err := proto.Marshal(patch)
	if err != nil {
		panic(err)
	}
	// tampered copy: flip one bit of the value MAC of mutation 2 (E1). The C++ side must reject it.
	tam := proto.Clone(patch).(*waServerSync.SyncdPatch)
	tb := tam.Mutations[2].Record.Value.Blob
	tb[len(tb)-1] ^= 0x01
	tamBytes, _ := proto.Marshal(tam)

	// ---- key share and key request messages (E2E Message) ----
	ksMsg := &waE2E.Message{ProtocolMessage: &waE2E.ProtocolMessage{
		Type: waE2E.ProtocolMessage_APP_STATE_SYNC_KEY_SHARE.Enum(),
		AppStateSyncKeyShare: &waE2E.AppStateSyncKeyShare{Keys: []*waE2E.AppStateSyncKey{{
			KeyID: &waE2E.AppStateSyncKeyId{KeyID: keyID},
			KeyData: &waE2E.AppStateSyncKeyData{
				KeyData:     keyData,
				Fingerprint: &waE2E.AppStateSyncKeyFingerprint{RawID: proto.Uint32(7), CurrentIndex: proto.Uint32(1), DeviceIndexes: []uint32{0, 3}},
				Timestamp:   proto.Int64(1767225600123),
			},
		}}},
	}}
	ksBytes, _ := proto.Marshal(ksMsg)
	krMsg := &waE2E.Message{ProtocolMessage: &waE2E.ProtocolMessage{
		Type: waE2E.ProtocolMessage_APP_STATE_SYNC_KEY_REQUEST.Enum(),
		AppStateSyncKeyRequest: &waE2E.AppStateSyncKeyRequest{KeyIDs: []*waE2E.AppStateSyncKeyId{{KeyID: keyID}}},
	}}
	krBytes, _ := proto.Marshal(krMsg)

	// final hash must equal a sequential build from the snapshot values
	out := map[string]any{
		"collection":        collection,
		"key_id":            hex.EncodeToString(keyID),
		"key_data":          hex.EncodeToString(keyData),
		"expanded_keys":     hex.EncodeToString(k.index) + hex.EncodeToString(k.valueEnc) + hex.EncodeToString(k.valueMAC) + hex.EncodeToString(k.snapshotMAC) + hex.EncodeToString(k.patchMAC),
		"snapshot":          hex.EncodeToString(snapBytes),
		"snapshot_hash_hex": snapHash,
		"patch":             hex.EncodeToString(patchBytes),
		"patch_tampered":    hex.EncodeToString(tamBytes),
		"patch_hash_hex":    hex.EncodeToString(state.hash[:]),
		"patch_version":     state.version,
		"key_share_message": hex.EncodeToString(ksBytes),
		"key_request_message": hex.EncodeToString(krBytes),
		"expect_contacts_after_snapshot": map[string]string{
			"31611111111@s.whatsapp.net": "Nigel Dokter",
			"31622222222@s.whatsapp.net": "Ann",
			"31633333333@s.whatsapp.net": "Pieter Bos",
		},
		"expect_contacts_after_patch": map[string]string{
			"31611111111@s.whatsapp.net": "Nigel Dokter (work)",
			"31622222222@s.whatsapp.net": "Ann",
			"31644444444@s.whatsapp.net": "Joop Visser",
		},
		"expect_removed_after_patch": []string{"31633333333@s.whatsapp.net"},
		"expect_unnamed_skipped":     []string{"31655555555@s.whatsapp.net"},
	}
	enc := json.NewEncoder(os.Stdout)
	enc.SetIndent("", "  ")
	if err := enc.Encode(out); err != nil {
		panic(fmt.Sprint(err))
	}
}