git.giacomolli.uk - onechat/blob - whatsapp_import.go

   1 package main
   2
   3 import (
   4         "archive/zip"
   5         "bufio"
   6         "context"
   7         "encoding/json"
   8         "errors"
   9         "fmt"
  10         "io"
  11         "maps"
  12         "os"
  13         "path/filepath"
  14         "regexp"
  15         "slices"
  16         "strings"
  17         "time"
  18
  19         "go.mau.fi/whatsmeow/store/sqlstore"
  20         waLog "go.mau.fi/whatsmeow/util/log"
  21 )
  22
  23 const (
  24         // whatsappImportMark marks history in the first line of a batch.
  25         whatsappImportMark = " (imported)"
  26
  27         // whatsappImportBytes limits one export's uncompressed chat text.
  28         whatsappImportBytes = 64 << 20
  29 )
  30
  31 var (
  32         whatsappImportStart   = regexp.MustCompile(`^\[(\d{2}/\d{2}/\d{4}, \d{2}:\d{2}:\d{2})\] ([^\r\n]+?): ?(.*)$`)
  33         whatsappImportHeader  = regexp.MustCompile(`^\[WhatsApp\] (.*), (?:group|private chat) ([^ ,]+), [A-Za-z]+ (\d{4}-\d{2}-\d{2})$`)
  34         whatsappImportUnknown = regexp.MustCompile(`^\[WhatsApp\] (.*), [A-Za-z]+ (\d{4}-\d{2}-\d{2})$`)
  35 )
  36
  37 // whatsappImport runs the offline import command. It holds the same token lock
  38 // as the daemon and probe, even for a plan. The plan reads no credentials for a
  39 // model and makes no network call. Only a completed rebuild replaces the agent
  40 // directory, and unfinished subagents prevent either mode.
  41 func whatsappImport(ctx context.Context, args []string, out io.Writer) (err error) {
  42         planOnly := len(args) > 0 && args[0] == "-n"
  43         if planOnly {
  44                 args = args[1:]
  45         }
  46
  47         if len(args) == 0 {
  48                 return errors.New("usage: onechat import [-n] ZIP [ZIP...]")
  49         }
  50
  51         for _, arg := range args {
  52                 if strings.HasPrefix(arg, "-") {
  53                         return errors.New("usage: onechat import [-n] ZIP [ZIP...]")
  54                 }
  55         }
  56
  57         dir, err := dataDir()
  58         if err != nil {
  59                 return fmt.Errorf("find import data directory: %w", err)
  60         }
  61
  62         tokens, err := lockTokens(dir)
  63         if err != nil {
  64                 return fmt.Errorf("lock import tokens: %w", err)
  65         }
  66
  67         defer func() {
  68                 closeErr := tokens.Close()
  69                 if closeErr != nil {
  70                         err = errors.Join(err, fmt.Errorf("close import token lock: %w", closeErr))
  71                 }
  72         }()
  73
  74         root := filepath.Join(dir, "agent")
  75         err = rebuildIdle(root)
  76         if err != nil {
  77                 return fmt.Errorf("check import work: %w", err)
  78         }
  79
  80         old, err := rebuildReadLog(ctx, root)
  81         if err != nil {
  82                 return fmt.Errorf("read import log: %w", err)
  83         }
  84
  85         names, owner, err := whatsappImportNames(ctx, dir)
  86         if err != nil {
  87                 return fmt.Errorf("read import names: %w", err)
  88         }
  89
  90         plan, err := whatsappImportPlan(ctx, args, old, nil, names, owner, time.Local)
  91         if err != nil {
  92                 return fmt.Errorf("plan import: %w", err)
  93         }
  94
  95         err = plan.Write(out)
  96         if err != nil {
  97                 return fmt.Errorf("write import plan: %w", err)
  98         }
  99
 100         if planOnly || len(plan.added) == 0 {
 101                 return nil
 102         }
 103
 104         // An existing agent is required so that a failed import never creates
 105         // or changes the original agent's configuration.
 106         _, err = os.Stat(filepath.Join(root, "config.json"))
 107         if err != nil {
 108                 return fmt.Errorf("read agent config: %w", err)
 109         }
 110
 111         agent, err := openAgent(root)
 112         if err != nil {
 113                 return fmt.Errorf("open import agent: %w", err)
 114         }
 115
 116         config, err := loadConfig(dir)
 117         if err != nil {
 118                 return fmt.Errorf("load import configuration: %w", err)
 119         }
 120
 121         address, auth, err := modelEndpoint(config, dir)
 122         if err != nil {
 123                 return fmt.Errorf("select import endpoint: %w", err)
 124         }
 125
 126         usage, err := os.CreateTemp(dir, ".import-usage-*.jsonl")
 127         if err != nil {
 128                 return fmt.Errorf("create import usage: %w", err)
 129         }
 130
 131         usagePath := usage.Name()
 132         defer func() {
 133                 closeErr := os.Remove(usagePath)
 134                 if closeErr != nil {
 135                         err = errors.Join(err, fmt.Errorf("remove import usage: %w", closeErr))
 136                 }
 137         }()
 138
 139         err = usage.Close()
 140         if err != nil {
 141                 return fmt.Errorf("close import usage: %w", err)
 142         }
 143
 144         client := newModelClient(address, auth, agent.CacheKey(), usagePath)
 145         backup, err := rebuild(ctx, agent, plan.added, client, newCompactor(client, agent.Prompt, config.Compactor, agent.Tools()), out)
 146         if err != nil {
 147                 return fmt.Errorf("rebuild imported memory: %w", err)
 148         }
 149
 150         _, err = fmt.Fprintf(out, "import complete: backup %s\n", backup)
 151         if err != nil {
 152                 return fmt.Errorf("write import result: %w", err)
 153         }
 154
 155         return nil
 156 }
 157
 158 // WhatsAppImportPlan holds the new chat-day batches and the size of the merged
 159 // log. Its estimate uses the measured cost of 1.3 calls per message, not a
 160 // promise about a particular export or compaction retry.
 161 type WhatsAppImportPlan struct {
 162         cutoffs    map[string]time.Time
 163         added      []DatedMessage
 164         chats      []string
 165         cutoff     time.Time
 166         first      time.Time
 167         last       time.Time
 168         bytes      int
 169         messages   int
 170         skipped    int
 171         ownerKnown bool
 172 }
 173
 174 // Write prints the plan without chat text or the owner's profile name.
 175 func (p WhatsAppImportPlan) Write(out io.Writer) error {
 176         cutoff := "none (no live WhatsApp batch)"
 177         if !p.cutoff.IsZero() {
 178                 cutoff = p.cutoff.Format(time.RFC3339)
 179         }
 180
 181         dateRange := "none"
 182         if !p.first.IsZero() {
 183                 dateRange = p.first.Format(time.RFC3339) + " to " + p.last.Format(time.RFC3339)
 184         }
 185
 186         owner := "known (mapped to You)"
 187         if !p.ownerKnown {
 188                 owner = "unknown (sender names kept as written)"
 189         }
 190
 191         _, err := fmt.Fprintf(out, "chats: %d\n", len(p.chats))
 192         if err != nil {
 193                 return fmt.Errorf("write import plan: %w", err)
 194         }
 195
 196         for _, chat := range p.chats {
 197                 line := "  " + chat
 198                 cutoff, found := p.cutoffs[chat]
 199                 if found {
 200                         line += " (cutoff: " + cutoff.Format(time.RFC3339) + ")"
 201                 }
 202
 203                 _, err = fmt.Fprintln(out, line)
 204                 if err != nil {
 205                         return fmt.Errorf("write import plan: %w", err)
 206                 }
 207         }
 208
 209         estimate := 1.3 * float64(p.messages)
 210         if len(p.added) == 0 {
 211                 estimate = 0
 212         }
 213
 214         lines := []string{
 215                 fmt.Sprintf("batches: %d (%d chat-days skipped)", len(p.added), p.skipped),
 216                 "dates: " + dateRange,
 217                 "cutoff: " + cutoff,
 218                 "owner: " + owner,
 219                 fmt.Sprintf("bytes: %d (new batch text)", p.bytes),
 220                 fmt.Sprintf("messages: %d (merged log)", p.messages),
 221                 fmt.Sprintf("estimated calls: %.0f (about 1.3 per message, limit %d)", estimate, rebuildCallLimit),
 222         }
 223
 224         _, err = io.WriteString(out, strings.Join(lines, "\n")+"\n")
 225         if err != nil {
 226                 return fmt.Errorf("write import plan: %w", err)
 227         }
 228
 229         return nil
 230 }
 231
 232 // whatsappImportNames reads only local session data. Group names that the
 233 // library keeps only in memory can also come from logged or pending batches.
 234 // Ambiguous names have no ID. A missing session keeps sender names unchanged.
 235 func whatsappImportNames(ctx context.Context, dir string) (names map[string]WhatsAppChat, owner string, err error) {
 236         names = make(map[string]WhatsAppChat)
 237         path := filepath.Join(dir, whatsappDirName, whatsappSessionFileName)
 238         _, err = os.Stat(path)
 239         if errors.Is(err, os.ErrNotExist) {
 240                 return names, "", nil
 241         }
 242
 243         if err != nil {
 244                 return nil, "", fmt.Errorf("read import session: %w", err)
 245         }
 246
 247         address, err := whatsappSessionAddress(path)
 248         if err != nil {
 249                 return nil, "", fmt.Errorf("find import session address: %w", err)
 250         }
 251
 252         container, err := sqlstore.New(ctx, "sqlite", address, waLog.Noop)
 253         if err != nil {
 254                 return nil, "", fmt.Errorf("open import session: %w", err)
 255         }
 256
 257         defer func() {
 258                 closeErr := container.Close()
 259                 if closeErr != nil {
 260                         err = errors.Join(err, fmt.Errorf("close import session: %w", closeErr))
 261                 }
 262         }()
 263
 264         device, err := container.GetFirstDevice(ctx)
 265         if err != nil {
 266                 return nil, "", fmt.Errorf("read import device: %w", err)
 267         }
 268
 269         if device.ID == nil {
 270                 return names, "", nil
 271         }
 272
 273         owner = device.PushName
 274         own := whatsappOwnIDs(device.ID.User, device.LID.User)
 275         contacts, err := device.Contacts.GetAllContacts(ctx)
 276         if err != nil {
 277                 return nil, "", fmt.Errorf("read import contacts: %w", err)
 278         }
 279
 280         for jid, contact := range contacts {
 281                 if slices.Contains(own, jid.User) {
 282                         continue
 283                 }
 284
 285                 for _, name := range []string{contact.FirstName, contact.FullName, contact.PushName, contact.BusinessName} {
 286                         if name != "" {
 287                                 whatsappImportName(names, WhatsAppChat{
 288                                         ID:    jid.ToNonAD().String(),
 289                                         Name:  name,
 290                                         Group: jid.Server == "g.us",
 291                                 })
 292                         }
 293                 }
 294         }
 295
 296         data, err := os.ReadFile(filepath.Join(dir, whatsappDirName, whatsappStateFileName))
 297         if errors.Is(err, os.ErrNotExist) {
 298                 return names, owner, nil
 299         }
 300
 301         if err != nil {
 302                 return nil, "", fmt.Errorf("read pending import names: %w", err)
 303         }
 304
 305         var state WhatsAppStateFile
 306
 307         err = json.Unmarshal(data, &state)
 308         if err != nil {
 309                 return nil, "", fmt.Errorf("decode pending import names: %w", err)
 310         }
 311
 312         for _, batch := range state.Batches {
 313                 whatsappImportName(names, batch.Chat)
 314         }
 315
 316         return names, owner, nil
 317 }
 318
 319 // whatsappImportName adds an exact match unless distinct IDs share its name.
 320 func whatsappImportName(names map[string]WhatsAppChat, chat WhatsAppChat) {
 321         old, found := names[chat.Name]
 322         if found && old.ID != chat.ID {
 323                 chat.ID = ""
 324         }
 325
 326         names[chat.Name] = chat
 327 }
 328
 329 // whatsappImportPlan filters each chat at its earliest live message across the
 330 // target and root logs. Only imported days in the target are skipped. Each
 331 // remaining chat-day becomes a work message dated at its last exported message.
 332 // Overlapping input chat-days are refused.
 333 func whatsappImportPlan(ctx context.Context, paths []string, old, rootLog []DatedMessage, names map[string]WhatsAppChat, owner string, location *time.Location) (WhatsAppImportPlan, error) {
 334         plan := WhatsAppImportPlan{
 335                 ownerKnown: owner != "",
 336                 cutoffs:    make(map[string]time.Time),
 337         }
 338
 339         names = maps.Clone(names)
 340         if names == nil {
 341                 names = make(map[string]WhatsAppChat)
 342         }
 343
 344         imported, err := whatsappImportHistory(old, rootLog, names, plan.cutoffs, location)
 345         if err != nil {
 346                 return plan, fmt.Errorf("read import history: %w", err)
 347         }
 348
 349         seen := make(map[string]bool)
 350         for _, path := range paths {
 351                 if ctx.Err() != nil {
 352                         return plan, ctx.Err()
 353                 }
 354
 355                 chat, messages, err := whatsappImportRead(path, owner, location)
 356                 if err != nil {
 357                         return plan, fmt.Errorf("read import export: %w", err)
 358                 }
 359
 360                 known, found := names[chat.Name]
 361                 if found {
 362                         chat = known
 363                 }
 364
 365                 if !slices.Contains(plan.chats, chat.Name) {
 366                         plan.chats = append(plan.chats, chat.Name)
 367                 }
 368
 369                 cutoff := plan.cutoffs[chat.Name]
 370                 if !cutoff.IsZero() && (plan.cutoff.IsZero() || cutoff.Before(plan.cutoff)) {
 371                         plan.cutoff = cutoff
 372                 }
 373
 374                 days := make(map[string][]WhatsAppMessage)
 375                 for _, message := range messages {
 376                         if !cutoff.IsZero() && !message.Time.Before(cutoff) {
 377                                 continue
 378                         }
 379
 380                         day := message.Time.In(location).Format(time.DateOnly)
 381                         days[day] = append(days[day], message)
 382                 }
 383
 384                 for _, day := range slices.Sorted(maps.Keys(days)) {
 385                         key := chat.Name + "\n" + day
 386                         if imported[key] {
 387                                 plan.skipped++
 388                                 continue
 389                         }
 390
 391                         if seen[key] {
 392                                 return plan, errors.New("import inputs overlap: give only one export for each chat and day")
 393                         }
 394
 395                         seen[key] = true
 396                         batch := WhatsAppBatch{
 397                                 Chat:     chat,
 398                                 Messages: days[day],
 399                         }
 400
 401                         text := batch.Text(location)
 402                         first, body, _ := strings.Cut(text, "\n")
 403                         if chat.ID == "" {
 404                                 first = "[WhatsApp] " + chat.Name + ", " + batch.Messages[0].Time.In(location).Format(whatsappBatchDayLayout)
 405                         }
 406
 407                         text = first + whatsappImportMark + "\n" + body
 408                         plan.added = append(plan.added, DatedMessage{
 409                                 Kind: "work",
 410                                 Text: text,
 411                                 Date: batch.Messages[len(batch.Messages)-1].Time.Format(time.RFC3339Nano),
 412                         })
 413
 414                         plan.bytes += len(text)
 415                         firstDate, lastDate := batch.Messages[0].Time, batch.Messages[len(batch.Messages)-1].Time
 416                         if plan.first.IsZero() || firstDate.Before(plan.first) {
 417                                 plan.first = firstDate
 418                         }
 419
 420                         if plan.last.IsZero() || lastDate.After(plan.last) {
 421                                 plan.last = lastDate
 422                         }
 423                 }
 424         }
 425
 426         plan.added, err = rebuildMerge(nil, plan.added)
 427         if err != nil {
 428                 return plan, fmt.Errorf("order import batches: %w", err)
 429         }
 430
 431         plan.messages = len(old) + len(plan.added)
 432         slices.Sort(plan.chats)
 433
 434         return plan, nil
 435 }
 436
 437 // whatsappImportHistory reads batch names and live cutoffs from both logs.
 438 // Imported chat-days belong only to the target log, which is first.
 439 func whatsappImportHistory(target, root []DatedMessage, names map[string]WhatsAppChat, cutoffs map[string]time.Time, location *time.Location) (map[string]bool, error) {
 440         imported := make(map[string]bool)
 441         for source, entries := range [][]DatedMessage{target, root} {
 442                 for _, entry := range entries {
 443                         if entry.Kind != "work" || !strings.HasPrefix(entry.Text, "[WhatsApp] ") {
 444                                 continue
 445                         }
 446
 447                         first, _, _ := strings.Cut(entry.Text, "\n")
 448                         marked := strings.HasSuffix(first, whatsappImportMark)
 449                         first = strings.TrimSuffix(first, whatsappImportMark)
 450                         match := whatsappImportHeader.FindStringSubmatch(first)
 451                         name, day := "", ""
 452                         if len(match) > 0 {
 453                                 name, day = match[1], match[3]
 454                                 whatsappImportName(names, WhatsAppChat{
 455                                         ID:    match[2],
 456                                         Name:  name,
 457                                         Group: strings.Contains(first, ", group "+match[2]+","),
 458                                 })
 459                         } else if marked {
 460                                 match = whatsappImportUnknown.FindStringSubmatch(first)
 461                                 if len(match) > 0 {
 462                                         name, day = match[1], match[2]
 463                                 }
 464                         }
 465
 466                         if name == "" {
 467                                 return nil, errors.New("invalid WhatsApp batch header")
 468                         }
 469
 470                         if marked {
 471                                 if source == 0 {
 472                                         imported[name+"\n"+day] = true
 473                                 }
 474
 475                                 continue
 476                         }
 477
 478                         _, body, found := strings.Cut(entry.Text, "\n")
 479                         if !found || len(body) < len("15:04 ") {
 480                                 return nil, errors.New("missing first WhatsApp message time")
 481                         }
 482
 483                         date, err := time.ParseInLocation("2006-01-02 15:04", day+" "+body[:5], location)
 484                         if err != nil {
 485                                 return nil, fmt.Errorf("parse first live message time: %w", err)
 486                         }
 487
 488                         if cutoffs[name].IsZero() || date.Before(cutoffs[name]) {
 489                                 cutoffs[name] = date
 490                         }
 491                 }
 492         }
 493
 494         return imported, nil
 495 }
 496
 497 // whatsappImportRead reads the single _chat.txt in an export without extracting
 498 // archive paths. Non-start lines continue the preceding message, and all
 499 // invisible left-to-right marks are removed. Dates use location. Media and
 500 // system messages remain text, including the written sender.
 501 func whatsappImportRead(path, owner string, location *time.Location) (chat WhatsAppChat, messages []WhatsAppMessage, err error) {
 502         chat.Name = strings.TrimPrefix(strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)), "WhatsApp Chat - ")
 503         if chat.Name == "" || strings.ContainsAny(chat.Name, "\r\n") {
 504                 return chat, nil, errors.New("read export: chat name must be one nonempty line")
 505         }
 506
 507         archive, err := zip.OpenReader(path)
 508         if err != nil {
 509                 return chat, nil, fmt.Errorf("open export: %w", err)
 510         }
 511
 512         defer func() {
 513                 closeErr := archive.Close()
 514                 if closeErr != nil {
 515                         err = errors.Join(err, fmt.Errorf("close export: %w", closeErr))
 516                 }
 517         }()
 518
 519         var file *zip.File
 520
 521         for _, entry := range archive.File {
 522                 if filepath.Base(entry.Name) == "_chat.txt" && !entry.FileInfo().IsDir() {
 523                         if file != nil {
 524                                 return chat, nil, errors.New("read export: more than one _chat.txt")
 525                         }
 526
 527                         file = entry
 528                 }
 529         }
 530
 531         if file == nil {
 532                 return chat, nil, errors.New("read export: _chat.txt is missing")
 533         }
 534
 535         if file.UncompressedSize64 > whatsappImportBytes {
 536                 return chat, nil, errors.New("read export: chat text exceeds 64 MiB")
 537         }
 538
 539         reader, err := file.Open()
 540         if err != nil {
 541                 return chat, nil, fmt.Errorf("open chat text: %w", err)
 542         }
 543
 544         defer func() {
 545                 closeErr := reader.Close()
 546                 if closeErr != nil {
 547                         err = errors.Join(err, fmt.Errorf("close export chat text: %w", closeErr))
 548                 }
 549         }()
 550
 551         scanner := bufio.NewScanner(io.LimitReader(reader, whatsappImportBytes+1))
 552         scanner.Buffer(make([]byte, 64*1024), buildLineBytes)
 553
 554         for scanner.Scan() {
 555                 line := strings.ReplaceAll(scanner.Text(), "\u200e", "")
 556                 match := whatsappImportStart.FindStringSubmatch(line)
 557                 if match == nil {
 558                         if len(messages) == 0 {
 559                                 if line == "" {
 560                                         continue
 561                                 }
 562
 563                                 return chat, nil, errors.New("read export: text before the first dated message")
 564                         }
 565
 566                         messages[len(messages)-1].Text += "\n" + line
 567
 568                         continue
 569                 }
 570
 571                 var date time.Time
 572
 573                 date, err = time.ParseInLocation("02/01/2006, 15:04:05", match[1], location)
 574                 if err != nil {
 575                         return chat, nil, fmt.Errorf("parse export local date: %w", err)
 576                 }
 577
 578                 messages = append(messages, WhatsAppMessage{
 579                         Sender: match[2],
 580                         Text:   match[3],
 581                         Time:   date,
 582                         Mine:   owner != "" && match[2] == owner && match[2] != chat.Name,
 583                 })
 584
 585                 if match[2] == chat.Name {
 586                         chat.Group = true
 587                 }
 588         }
 589
 590         if scanner.Err() != nil {
 591                 return chat, nil, fmt.Errorf("read chat text: %w", scanner.Err())
 592         }
 593
 594         if len(messages) == 0 {
 595                 return chat, nil, errors.New("read export: no dated messages")
 596         }
 597
 598         slices.SortStableFunc(messages, func(a, b WhatsAppMessage) int {
 599                 return a.Time.Compare(b.Time)
 600         })
 601
 602         return chat, messages, nil
 603 }