wx-win-agent/node-agent/WxAgent.Windows/WechatMessageDbReader.cs

using System.Globalization;
using WxAgent.Core;

namespace WxAgent.Windows;

/// <summary>Reads chat messages from the account's SQLCipher message databases using cached verified keys.</summary>
public static class WechatMessageDbReader
{
    public static async Task<IReadOnlyList<DbMessage>> ReadAsync(
        string accountRootPath, IReadOnlyList<DatabaseKeyEvidence> databases, string chatId, int limit, CancellationToken cancellationToken, long? localId = null)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(accountRootPath);
        ArgumentException.ThrowIfNullOrWhiteSpace(chatId);
        if (limit is < 1 or > 500) throw new WxAgentException(WxAgentErrorCode.InvalidArgument, "Message limit must be between 1 and 500.");

        // AccountRootPath points at the account's db_storage folder; the directory above it is <wxid>_<random>.
        var accountRootDirectoryName = Path.GetFileName(Path.GetDirectoryName(Path.TrimEndingDirectorySeparator(Path.GetFullPath(accountRootPath))));

        var tableName = WechatDbMessage.TableNameFor(chatId);
        if (!tableName.StartsWith("Msg_", StringComparison.Ordinal) || tableName.Length != 36 || !tableName[4..].All(Uri.IsHexDigit))
        {
            throw new WxAgentException(WxAgentErrorCode.InvalidArgument, "Message table name could not be derived from the chat id.");
        }
        var messageDatabase = await FindMessageDatabaseAsync(accountRootPath, databases, tableName, cancellationToken).ConfigureAwait(false);
        if (messageDatabase is null)
        {
            return Array.Empty<DbMessage>();
        }

        var sql = "SELECT m.local_id, m.server_id, m.local_type, m.create_time, " +
            "hex(m.message_content) AS hex_content, m.WCDB_CT_message_content AS is_compressed, n.user_name AS sender_wxid " +
            $"FROM \"{tableName}\" m LEFT JOIN Name2Id n ON m.real_sender_id = n.rowid " +
            "WHERE ($localId IS NULL OR m.local_id = $localId) " +
            "ORDER BY m.create_time DESC, m.local_id DESC LIMIT $limit;";
        var rows = await SqlCipherDatabaseReader.QueryRowsAsync(messageDatabase.Value.Path, messageDatabase.Value.EncKey, sql,
            [new KeyValuePair<string, object?>("$limit", (long)limit), new KeyValuePair<string, object?>("$localId", localId)], cancellationToken).ConfigureAwait(false);

        var senders = rows.Select(row => row.GetValueOrDefault("sender_wxid") as string)
            .Where(sender => !string.IsNullOrEmpty(sender)).Cast<string>().Distinct(StringComparer.Ordinal).ToArray();
        var contacts = await ReadContactsAsync(accountRootPath, databases, senders, cancellationToken).ConfigureAwait(false);

        var messages = new List<DbMessage>(rows.Count);
        foreach (var row in rows)
        {
            var senderWxId = row.GetValueOrDefault("sender_wxid") as string;
            var content = WechatDbMessage.DecodeContent(
                row.GetValueOrDefault("hex_content") as string ?? string.Empty,
                Convert.ToInt64(row.GetValueOrDefault("is_compressed") ?? 0L, CultureInfo.InvariantCulture) != 0);
            var contact = senderWxId is null ? null : contacts.GetValueOrDefault(senderWxId);
            messages.Add(new DbMessage(
                Convert.ToInt64(row.GetValueOrDefault("local_id") ?? 0L, CultureInfo.InvariantCulture),
                Convert.ToInt64(row.GetValueOrDefault("server_id") ?? 0L, CultureInfo.InvariantCulture),
                chatId,
                senderWxId,
                contact?.Name,
                contact?.AvatarUrl,
                Convert.ToInt64(row.GetValueOrDefault("local_type") ?? 0L, CultureInfo.InvariantCulture),
                content,
                DateTimeOffset.FromUnixTimeSeconds(Convert.ToInt64(row.GetValueOrDefault("create_time") ?? 0L, CultureInfo.InvariantCulture)),
                WechatDbMessage.IsSelf(senderWxId, accountRootDirectoryName)));
        }

        return messages;
    }

    public static async Task<DbMergedChat> ReadMergedAsync(string accountRootPath, IReadOnlyList<DatabaseKeyEvidence> databases,
        string chatId, long localId, CancellationToken cancellationToken, string? databaseRelativePath = null)
    {
        if (localId <= 0) throw new WxAgentException(WxAgentErrorCode.InvalidArgument, "A positive local message ID is required.");
        var candidates = databases.Where(database => database.RelativePath.StartsWith("message/", StringComparison.OrdinalIgnoreCase)
            && (databaseRelativePath is null || string.Equals(database.RelativePath, databaseRelativePath.Replace('\\', '/'), StringComparison.OrdinalIgnoreCase))).ToArray();
        if (candidates.Length == 0) throw new WxAgentException(WxAgentErrorCode.DatabaseKeyNotFound, "No matching message database has a cached key.");
        (string Path, DbMessage Message)? found = null;
        foreach (var database in candidates)
        {
            cancellationToken.ThrowIfCancellationRequested();
            var selected = databases.Where(entry => !entry.RelativePath.StartsWith("message/", StringComparison.OrdinalIgnoreCase) || entry == database).ToArray();
            var rows = await ReadAsync(accountRootPath, selected, chatId, 2, cancellationToken, localId).ConfigureAwait(false);
            if (rows.Count == 0) continue;
            if (rows.Count != 1 || found is not null)
                throw new WxAgentException(WxAgentErrorCode.InvalidOperationState, "The local ID is ambiguous across message shards; specify --database.");
            found = (database.RelativePath, rows[0]);
        }
        if (found is null) throw new WxAgentException(WxAgentErrorCode.ControlNotFound, "The selected database message was not found.");
        try { return new DbMergedChat(found.Value.Path, found.Value.Message, WechatMergedChatParser.Parse(found.Value.Message.Content)); }
        catch (InvalidDataException ex) { throw new WxAgentException(WxAgentErrorCode.DatabaseOpenFailed, ex.Message, ex); }
    }

    private static async Task<IReadOnlyDictionary<string, WechatContactName>> ReadContactsAsync(
        string accountRootPath, IReadOnlyList<DatabaseKeyEvidence> databases, IReadOnlyList<string> senders, CancellationToken cancellationToken)
    {
        var result = new Dictionary<string, WechatContactName>(StringComparer.Ordinal);
        if (senders.Count == 0) return result;

        var contactDatabase = databases.FirstOrDefault(database => string.Equals(database.RelativePath, "contact/contact.db", StringComparison.OrdinalIgnoreCase));
        if (contactDatabase is null) return result;
        var path = Path.GetFullPath(Path.Combine(accountRootPath, "contact", "contact.db"));

        var placeholders = string.Join(", ", Enumerable.Range(0, senders.Count).Select(index => $"$p{index}"));
        var parameters = senders.Select((sender, index) => new KeyValuePair<string, object?>($"$p{index}", sender)).ToArray();
        var sql = $"SELECT username, nick_name, remark, small_head_url FROM contact WHERE username IN ({placeholders});";
        var rows = await SqlCipherDatabaseReader.QueryRowsAsync(path, contactDatabase.EncKey, sql, parameters, cancellationToken).ConfigureAwait(false);
        foreach (var row in rows)
        {
            var username = row.GetValueOrDefault("username") as string;
            if (string.IsNullOrEmpty(username)) continue;
            var remark = row.GetValueOrDefault("remark") as string;
            var nickname = row.GetValueOrDefault("nick_name") as string;
            var name = !string.IsNullOrEmpty(remark) ? remark : !string.IsNullOrEmpty(nickname) ? nickname : username;
            result[username] = new WechatContactName(name, row.GetValueOrDefault("small_head_url") as string);
        }

        return result;
    }

    private static async Task<(string Path, string EncKey)?> FindMessageDatabaseAsync(
        string accountRootPath, IReadOnlyList<DatabaseKeyEvidence> databases, string tableName, CancellationToken cancellationToken)
    {
        var candidates = databases
            .Where(database => database.RelativePath.StartsWith("message/", StringComparison.OrdinalIgnoreCase)
                && (database.RelativePath.EndsWith(".db", StringComparison.OrdinalIgnoreCase)))
            .OrderBy(database => database.RelativePath, StringComparer.OrdinalIgnoreCase)
            .ToArray();
        foreach (var database in candidates)
        {
            cancellationToken.ThrowIfCancellationRequested();
            var root = Path.GetFullPath(accountRootPath).TrimEnd(Path.DirectorySeparatorChar, Path.AltDirectorySeparatorChar) + Path.DirectorySeparatorChar;
            var relative = database.RelativePath.Replace('/', Path.DirectorySeparatorChar);
            var path = Path.GetFullPath(Path.Combine(root, relative));
            if (Path.IsPathRooted(relative) || !path.StartsWith(root, StringComparison.OrdinalIgnoreCase))
                throw new WxAgentException(WxAgentErrorCode.InvalidArgument, "The database path must remain inside the selected account root.");
            var sql = "SELECT name FROM sqlite_master WHERE type = 'table' AND name = $name;";
            var rows = await SqlCipherDatabaseReader.QueryRowsAsync(path, database.EncKey, sql,
                [new KeyValuePair<string, object?>("$name", tableName)], cancellationToken).ConfigureAwait(false);
            if (rows.Count > 0) return (path, database.EncKey);
        }

        return null;
    }

    private sealed record WechatContactName(string Name, string? AvatarUrl);
}