InfiniteAPI/src/Utils/tc-token-utils.ts

import type { SignalKeyStoreWithTransaction } from '../Types'
import type { BinaryNode } from '../WABinary'
import {
	getBinaryNodeChild,
	getBinaryNodeChildren,
	isAnyLidUser,
	isAnyPnUser,
	isHostedLidUser,
	isHostedPnUser,
	isJidMetaAI,
	isLidUser,
	isPnUser,
	jidNormalizedUser
} from '../WABinary'

/** 7 days in seconds — matches WA Web AB prop tctoken_duration */
const TC_TOKEN_BUCKET_DURATION = 604800
/** 4 buckets → ~28-day rolling window — matches WA Web AB prop tctoken_num_buckets */
const TC_TOKEN_NUM_BUCKETS = 4

/**
 * Sentinel key under the `tctoken` store holding a JSON array of tracked storage JIDs
 * for cross-session pruning. Mirrors WA Web's CLEAN_TC_TOKENS index lookup.
 *
 * Exported so other modules (messages-recv, messages-send, process-message) reference
 * the same constant. The fork previously inlined this string in messages-recv.ts;
 * keep the value identical (`'__index'`) for backward compatibility with persisted state.
 */
export const TC_TOKEN_INDEX_KEY = '__index'

// Phone-number pattern matching WABinary's isJidBot, applied against the user part so
// the check is invariant to @c.us ↔ @s.whatsapp.net normalization.
const BOT_PHONE_REGEX = /^1313555\d{4}$|^131655500\d{2}$/

/**
 * Mirrors WA Web's `Wid.isRegularUser()` (user ∧ ¬PSA ∧ ¬Bot). Used to gate tctoken
 * storage against malformed notifications — WA Web filters server-side but we
 * defend here for parity with `WAWebSetTcTokenChatAction.handleIncomingTcToken`.
 * Works for both pre- and post-normalized JIDs (`@c.us` vs `@s.whatsapp.net`).
 */
export function isRegularUser(jid: string | undefined): boolean {
	if (!jid) return false
	const user = jid.split('@')[0] ?? ''
	if (user === '0') return false // PSA
	if (BOT_PHONE_REGEX.test(user)) return false // Bot by phone pattern
	if (isJidMetaAI(jid)) return false // MetaAI (@bot server)
	return !!(isPnUser(jid) || isLidUser(jid) || isHostedPnUser(jid) || isHostedLidUser(jid) || jid.endsWith('@c.us'))
}

/** Read the persisted tctoken JID index and return its entries (never contains the sentinel key itself). */
export async function readTcTokenIndex(keys: SignalKeyStoreWithTransaction): Promise<string[]> {
	const data = await keys.get('tctoken', [TC_TOKEN_INDEX_KEY])
	const entry = data[TC_TOKEN_INDEX_KEY]
	if (!entry?.token?.length) return []
	try {
		const parsed = JSON.parse(Buffer.from(entry.token).toString())
		if (!Array.isArray(parsed)) return []
		return parsed.filter((j): j is string => typeof j === 'string' && j.length > 0 && j !== TC_TOKEN_INDEX_KEY)
	} catch {
		return []
	}
}

/**
 * Build a SignalDataSet fragment that writes the merged index (persisted ∪ added)
 * under the sentinel key. Lets callers update the index without clobbering writes
 * made by other layers (history sync, concurrent sessions on the same store).
 */
export async function buildMergedTcTokenIndexWrite(
	keys: SignalKeyStoreWithTransaction,
	addedJids: Iterable<string>
): Promise<{ [TC_TOKEN_INDEX_KEY]: { token: Buffer } }> {
	const persisted = await readTcTokenIndex(keys)
	const merged = new Set(persisted)
	for (const jid of addedJids) {
		if (jid && jid !== TC_TOKEN_INDEX_KEY) merged.add(jid)
	}

	return {
		[TC_TOKEN_INDEX_KEY]: { token: Buffer.from(JSON.stringify([...merged])) }
	}
}

/**
 * Check if a received token is expired using WA Web's rolling bucket algorithm.
 * Reference: WAWebTrustedContactsUtils.isTokenExpired
 *
 * Uses Receiver mode constants (tctoken_duration, tctoken_num_buckets).
 * NOTE: WA Web distinguishes Sender vs Receiver mode via AB props
 * (tctoken_duration_sender / tctoken_num_buckets_sender). Currently both
 * use identical values (604800 / 4), so we use a single function for both.
 * If WA ever diverges these, add a `mode` parameter here.
 */
export function isTcTokenExpired(timestamp: number | string | null | undefined): boolean {
	if (timestamp === null || timestamp === undefined) return true
	const ts = typeof timestamp === 'string' ? Number(timestamp) : timestamp
	if (isNaN(ts)) return true
	const now = Math.floor(Date.now() / 1000)
	const currentBucket = Math.floor(now / TC_TOKEN_BUCKET_DURATION)
	const cutoffBucket = currentBucket - (TC_TOKEN_NUM_BUCKETS - 1)
	const cutoffTimestamp = cutoffBucket * TC_TOKEN_BUCKET_DURATION
	return ts < cutoffTimestamp
}

/**
 * Check if we should issue a new token to this contact (bucket boundary crossed).
 * Reference: WAWebTrustedContactsUtils.shouldSendNewToken
 *
 * Returns true if senderTimestamp is null/undefined or in a previous bucket.
 */
export function shouldSendNewTcToken(senderTimestamp: number | undefined): boolean {
	if (senderTimestamp === undefined) return true
	const now = Math.floor(Date.now() / 1000)
	const currentBucket = Math.floor(now / TC_TOKEN_BUCKET_DURATION)
	const senderBucket = Math.floor(senderTimestamp / TC_TOKEN_BUCKET_DURATION)
	return currentBucket > senderBucket
}

/**
 * Resolve a JID to its LID for tctoken storage, mirroring how Signal sessions
 * use LID keys via resolveLIDSignalAddress.
 *
 * WA Web always resolves to LID before storing/looking up tctokens:
 * `senderLid ?? toLid(from)` (WAWebSetTcTokenChatAction.handleIncomingTcToken)
 *
 * @param jid - The JID to resolve (can be PN or LID)
 * @param getLIDForPN - Resolver function (from lidMapping)
 * @returns The LID if mapping exists, otherwise the original JID
 */
export async function resolveTcTokenJid(
	jid: string,
	getLIDForPN: (pn: string) => Promise<string | null>
): Promise<string> {
	const normalized = jidNormalizedUser(jid)
	if (isLidUser(normalized)) return normalized
	const lid = await getLIDForPN(normalized)
	return lid ?? normalized
}

/**
 * Resolve target JID for issuing privacy token based on AB prop 14303
 * (`lid_trusted_token_issue_to_lid`). When the prop is on, issuance goes to
 * the LID; when off, it goes to the PN. Returns the original JID if no
 * mapping is found in either direction.
 *
 * Normalizes the JID upfront and uses the `isAny*` helpers so callers can pass
 * `@c.us`, `@s.whatsapp.net`, `@hosted`, `@hosted.lid`, `@lid` or device-specific
 * forms — `LIDMappingStore.getLIDForPN` early-returns unless `isAnyPnUser`,
 * so unnormalized inputs would silently bypass routing.
 *
 * Reference: WAWebTrustedContactsManager.issuePrivacyTokens
 */
export async function resolveIssuanceJid(
	jid: string,
	issueToLid: boolean,
	getLIDForPN: (pn: string) => Promise<string | null>,
	getPNForLID?: (lid: string) => Promise<string | null>
): Promise<string> {
	const normalized = jidNormalizedUser(jid)

	if (issueToLid) {
		if (isAnyLidUser(normalized)) return normalized
		if (!isAnyPnUser(normalized)) return normalized
		const lid = await getLIDForPN(normalized)
		return lid ?? normalized
	}

	if (!isAnyLidUser(normalized)) return normalized
	if (getPNForLID) {
		const pn = await getPNForLID(normalized)
		return pn ?? normalized
	}

	return normalized
}

type TcTokenParams = {
	jid: string
	baseContent?: BinaryNode[]
	authState: {
		keys: SignalKeyStoreWithTransaction
	}
	getLIDForPN?: (pn: string) => Promise<string | null>
}

export async function buildTcTokenFromJid({
	authState,
	jid,
	baseContent = [],
	getLIDForPN
}: TcTokenParams): Promise<BinaryNode[] | undefined> {
	try {
		const storageJid = getLIDForPN ? await resolveTcTokenJid(jid, getLIDForPN) : jid
		const tcTokenData = await authState.keys.get('tctoken', [storageJid])
		const entry = tcTokenData?.[storageJid]
		const tcTokenBuffer = entry?.token

		if (!tcTokenBuffer?.length || isTcTokenExpired(entry?.timestamp)) {
			// Opportunistic cleanup: remove expired token from store
			if (tcTokenBuffer) {
				await authState.keys.set({ tctoken: { [storageJid]: null } })
			}

			return baseContent.length > 0 ? baseContent : undefined
		}

		baseContent.push({
			tag: 'tctoken',
			attrs: {},
			content: tcTokenBuffer
		})

		return baseContent
	} catch (error) {
		return baseContent.length > 0 ? baseContent : undefined
	}
}

type StoreTcTokensParams = {
	result: BinaryNode
	fallbackJid: string
	keys: SignalKeyStoreWithTransaction
	getLIDForPN: (pn: string) => Promise<string | null>
	/** Optional callback when a new JID is stored (for index tracking) */
	onNewJidStored?: (jid: string) => void
}

/**
 * Parse and store tctoken(s) from an IQ result node.
 * Includes timestamp monotonicity guard matching WA Web's handleIncomingTcToken.
 * Used by both the blocking fetch (messages-send) and IQ response (messages-recv) paths.
 */
export async function storeTcTokensFromIqResult({
	result,
	fallbackJid,
	keys,
	getLIDForPN,
	onNewJidStored
}: StoreTcTokensParams) {
	const tokensNode = getBinaryNodeChild(result, 'tokens')
	if (!tokensNode) return

	const tokenNodes = getBinaryNodeChildren(tokensNode, 'token')
	for (const tokenNode of tokenNodes) {
		if (tokenNode.attrs.type !== 'trusted_contact' || !(tokenNode.content instanceof Uint8Array)) {
			continue
		}

		// In notifications, tokenNode.attrs.jid is OUR own device JID, not the sender's.
		// Prefer fallbackJid (resolved from notification's `from` / `sender_lid`) so the
		// token is stored under the peer's JID, never under self.
		const rawJid = jidNormalizedUser(fallbackJid || tokenNode.attrs.jid)
		// Defense against malformed notifications (PSA WID '0', bots, MetaAI). WA Web
		// filters these server-side; we mirror Wid.isRegularUser() locally.
		if (!isRegularUser(rawJid)) continue
		const storageJid = await resolveTcTokenJid(rawJid, getLIDForPN)
		const existingTcData = await keys.get('tctoken', [storageJid])
		const existingEntry = existingTcData[storageJid]

		// Timestamp monotonicity guard — only store if incoming timestamp >= existing
		// Matches WA Web handleIncomingTcToken
		const existingTs = existingEntry?.timestamp ? Number(existingEntry.timestamp) : 0
		const incomingTs = tokenNode.attrs.t ? Number(tokenNode.attrs.t) : 0
		if (existingTs > 0 && incomingTs > 0 && existingTs > incomingTs) {
			continue
		}

		// Don't store timestamp-less tokens at all — isTcTokenExpired treats them
		// as immediately expired regardless of whether an existing entry is present
		if (!incomingTs) {
			continue
		}

		const tokenEntry = {
			...existingEntry,
			token: Buffer.from(tokenNode.content),
			timestamp: tokenNode.attrs.t,
			// WABA Android: resets real_issue_timestamp to null when storing a new token
			// (UPDATE wa_trusted_contacts_send SET real_issue_timestamp=null)
			realIssueTimestamp: null
		}

		// Store under resolved storageJid AND under fallbackJid (PN) for reliable lookup
		// The read path may resolve to a different LID than the store path
		const normalizedFallback = jidNormalizedUser(fallbackJid)
		const keysToStore: Record<string, typeof tokenEntry | null> = {
			[storageJid]: tokenEntry
		}
		if (normalizedFallback !== storageJid) {
			keysToStore[normalizedFallback] = tokenEntry
		}

		await keys.set({ tctoken: keysToStore })
		onNewJidStored?.(storageJid)
	}
}