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, 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), 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         added      []DatedMessage
 163         chats      []string
 164         cutoff     time.Time
 165         first      time.Time
 166         last       time.Time
 167         bytes      int
 168         messages   int
 169         skipped    int
 170         ownerKnown bool
 171 }
 172
 173 // Write prints the plan without chat text or the owner's profile name.
 174 func (p WhatsAppImportPlan) Write(out io.Writer) error {
 175         cutoff := "none (no live WhatsApp batch)"
 176         if !p.cutoff.IsZero() {
 177                 cutoff = p.cutoff.Format(time.RFC3339)
 178         }
 179
 180         dateRange := "none"
 181         if !p.first.IsZero() {
 182                 dateRange = p.first.Format(time.RFC3339) + " to " + p.last.Format(time.RFC3339)
 183         }
 184
 185         owner := "known (mapped to You)"
 186         if !p.ownerKnown {
 187                 owner = "unknown (sender names kept as written)"
 188         }
 189
 190         _, err := fmt.Fprintf(out, "chats: %d\n", len(p.chats))
 191         if err != nil {
 192                 return fmt.Errorf("write import plan: %w", err)
 193         }
 194
 195         for _, chat := range p.chats {
 196                 _, err = fmt.Fprintf(out, "  %s\n", chat)
 197                 if err != nil {
 198                         return fmt.Errorf("write import plan: %w", err)
 199                 }
 200         }
 201
 202         estimate := 1.3 * float64(p.messages)
 203         if len(p.added) == 0 {
 204                 estimate = 0
 205         }
 206
 207         lines := []string{
 208                 fmt.Sprintf("batches: %d (%d chat-days skipped)", len(p.added), p.skipped),
 209                 "dates: " + dateRange,
 210                 "cutoff: " + cutoff,
 211                 "owner: " + owner,
 212                 fmt.Sprintf("bytes: %d (new batch text)", p.bytes),
 213                 fmt.Sprintf("messages: %d (merged log)", p.messages),
 214                 fmt.Sprintf("estimated calls: %.0f (about 1.3 per message, limit %d)", estimate, rebuildCallLimit),
 215         }
 216
 217         _, err = io.WriteString(out, strings.Join(lines, "\n")+"\n")
 218         if err != nil {
 219                 return fmt.Errorf("write import plan: %w", err)
 220         }
 221
 222         return nil
 223 }
 224
 225 // whatsappImportNames reads only local session data. Group names that the
 226 // library keeps only in memory can also come from logged or pending batches.
 227 // Ambiguous names have no ID. A missing session keeps sender names unchanged.
 228 func whatsappImportNames(ctx context.Context, dir string) (names map[string]WhatsAppChat, owner string, err error) {
 229         names = make(map[string]WhatsAppChat)
 230         path := filepath.Join(dir, whatsappDirName, whatsappSessionFileName)
 231         _, err = os.Stat(path)
 232         if errors.Is(err, os.ErrNotExist) {
 233                 return names, "", nil
 234         }
 235
 236         if err != nil {
 237                 return nil, "", fmt.Errorf("read import session: %w", err)
 238         }
 239
 240         address, err := whatsappSessionAddress(path)
 241         if err != nil {
 242                 return nil, "", fmt.Errorf("find import session address: %w", err)
 243         }
 244
 245         container, err := sqlstore.New(ctx, "sqlite", address, waLog.Noop)
 246         if err != nil {
 247                 return nil, "", fmt.Errorf("open import session: %w", err)
 248         }
 249
 250         defer func() {
 251                 closeErr := container.Close()
 252                 if closeErr != nil {
 253                         err = errors.Join(err, fmt.Errorf("close import session: %w", closeErr))
 254                 }
 255         }()
 256
 257         device, err := container.GetFirstDevice(ctx)
 258         if err != nil {
 259                 return nil, "", fmt.Errorf("read import device: %w", err)
 260         }
 261
 262         if device.ID == nil {
 263                 return names, "", nil
 264         }
 265
 266         owner = device.PushName
 267         own := whatsappOwnIDs(device.ID.User, device.LID.User)
 268         contacts, err := device.Contacts.GetAllContacts(ctx)
 269         if err != nil {
 270                 return nil, "", fmt.Errorf("read import contacts: %w", err)
 271         }
 272
 273         for jid, contact := range contacts {
 274                 if slices.Contains(own, jid.User) {
 275                         continue
 276                 }
 277
 278                 for _, name := range []string{contact.FirstName, contact.FullName, contact.PushName, contact.BusinessName} {
 279                         if name != "" {
 280                                 whatsappImportName(names, WhatsAppChat{
 281                                         ID:    jid.ToNonAD().String(),
 282                                         Name:  name,
 283                                         Group: jid.Server == "g.us",
 284                                 })
 285                         }
 286                 }
 287         }
 288
 289         data, err := os.ReadFile(filepath.Join(dir, whatsappDirName, whatsappStateFileName))
 290         if errors.Is(err, os.ErrNotExist) {
 291                 return names, owner, nil
 292         }
 293
 294         if err != nil {
 295                 return nil, "", fmt.Errorf("read pending import names: %w", err)
 296         }
 297
 298         var state WhatsAppStateFile
 299
 300         err = json.Unmarshal(data, &state)
 301         if err != nil {
 302                 return nil, "", fmt.Errorf("decode pending import names: %w", err)
 303         }
 304
 305         for _, batch := range state.Batches {
 306                 whatsappImportName(names, batch.Chat)
 307         }
 308
 309         return names, owner, nil
 310 }
 311
 312 // whatsappImportName adds an exact match unless distinct IDs share its name.
 313 func whatsappImportName(names map[string]WhatsAppChat, chat WhatsAppChat) {
 314         old, found := names[chat.Name]
 315         if found && old.ID != chat.ID {
 316                 chat.ID = ""
 317         }
 318
 319         names[chat.Name] = chat
 320 }
 321
 322 // whatsappImportPlan filters exports at the first message of the earliest live
 323 // batch and skips chat-days already imported. Each remaining chat-day becomes a
 324 // work message dated at its last exported message. Overlapping chat-days in the
 325 // inputs are refused, not silently combined or duplicated.
 326 func whatsappImportPlan(ctx context.Context, paths []string, old []DatedMessage, names map[string]WhatsAppChat, owner string, location *time.Location) (WhatsAppImportPlan, error) {
 327         plan := WhatsAppImportPlan{ownerKnown: owner != ""}
 328         names = maps.Clone(names)
 329         if names == nil {
 330                 names = make(map[string]WhatsAppChat)
 331         }
 332
 333         imported := make(map[string]bool)
 334         for _, entry := range old {
 335                 if entry.Kind != "work" || !strings.HasPrefix(entry.Text, "[WhatsApp] ") {
 336                         continue
 337                 }
 338
 339                 first, _, _ := strings.Cut(entry.Text, "\n")
 340                 marked := strings.HasSuffix(first, whatsappImportMark)
 341                 first = strings.TrimSuffix(first, whatsappImportMark)
 342                 match := whatsappImportHeader.FindStringSubmatch(first)
 343                 name, day := "", ""
 344                 if len(match) > 0 {
 345                         name, day = match[1], match[3]
 346                         whatsappImportName(names, WhatsAppChat{
 347                                 ID:    match[2],
 348                                 Name:  name,
 349                                 Group: strings.Contains(first, ", group "+match[2]+","),
 350                         })
 351                 } else if marked {
 352                         match = whatsappImportUnknown.FindStringSubmatch(first)
 353                         if len(match) > 0 {
 354                                 name, day = match[1], match[2]
 355                         }
 356                 }
 357
 358                 if name == "" {
 359                         return plan, errors.New("read WhatsApp history: invalid batch header")
 360                 }
 361
 362                 if marked {
 363                         imported[name+"\n"+day] = true
 364                         continue
 365                 }
 366
 367                 _, body, found := strings.Cut(entry.Text, "\n")
 368                 if !found || len(body) < len("15:04 ") {
 369                         return plan, errors.New("read WhatsApp history: missing first message time")
 370                 }
 371
 372                 date, err := time.ParseInLocation("2006-01-02 15:04", day+" "+body[:5], location)
 373                 if err != nil {
 374                         return plan, fmt.Errorf("parse first live message time: %w", err)
 375                 }
 376
 377                 if plan.cutoff.IsZero() || date.Before(plan.cutoff) {
 378                         plan.cutoff = date
 379                 }
 380         }
 381
 382         seen := make(map[string]bool)
 383         for _, path := range paths {
 384                 if ctx.Err() != nil {
 385                         return plan, ctx.Err()
 386                 }
 387
 388                 chat, messages, err := whatsappImportRead(path, owner, location)
 389                 if err != nil {
 390                         return plan, fmt.Errorf("read import export: %w", err)
 391                 }
 392
 393                 known, found := names[chat.Name]
 394                 if found {
 395                         chat = known
 396                 }
 397
 398                 if !slices.Contains(plan.chats, chat.Name) {
 399                         plan.chats = append(plan.chats, chat.Name)
 400                 }
 401
 402                 days := make(map[string][]WhatsAppMessage)
 403                 for _, message := range messages {
 404                         if !plan.cutoff.IsZero() && !message.Time.Before(plan.cutoff) {
 405                                 continue
 406                         }
 407
 408                         day := message.Time.In(location).Format(time.DateOnly)
 409                         days[day] = append(days[day], message)
 410                 }
 411
 412                 for _, day := range slices.Sorted(maps.Keys(days)) {
 413                         key := chat.Name + "\n" + day
 414                         if imported[key] {
 415                                 plan.skipped++
 416                                 continue
 417                         }
 418
 419                         if seen[key] {
 420                                 return plan, errors.New("import inputs overlap: give only one export for each chat and day")
 421                         }
 422
 423                         seen[key] = true
 424                         batch := WhatsAppBatch{
 425                                 Chat:     chat,
 426                                 Messages: days[day],
 427                         }
 428
 429                         text := batch.Text(location)
 430                         first, body, _ := strings.Cut(text, "\n")
 431                         if chat.ID == "" {
 432                                 first = "[WhatsApp] " + chat.Name + ", " + batch.Messages[0].Time.In(location).Format(whatsappBatchDayLayout)
 433                         }
 434
 435                         text = first + whatsappImportMark + "\n" + body
 436                         plan.added = append(plan.added, DatedMessage{
 437                                 Kind: "work",
 438                                 Text: text,
 439                                 Date: batch.Messages[len(batch.Messages)-1].Time.Format(time.RFC3339Nano),
 440                         })
 441
 442                         plan.bytes += len(text)
 443                         firstDate, lastDate := batch.Messages[0].Time, batch.Messages[len(batch.Messages)-1].Time
 444                         if plan.first.IsZero() || firstDate.Before(plan.first) {
 445                                 plan.first = firstDate
 446                         }
 447
 448                         if plan.last.IsZero() || lastDate.After(plan.last) {
 449                                 plan.last = lastDate
 450                         }
 451                 }
 452         }
 453
 454         var err error
 455
 456         plan.added, err = rebuildMerge(nil, plan.added)
 457         if err != nil {
 458                 return plan, fmt.Errorf("order import batches: %w", err)
 459         }
 460
 461         plan.messages = len(old) + len(plan.added)
 462         slices.Sort(plan.chats)
 463
 464         return plan, nil
 465 }
 466
 467 // whatsappImportRead reads the single _chat.txt in an export without extracting
 468 // archive paths. Non-start lines continue the preceding message, and all
 469 // invisible left-to-right marks are removed. Dates use location. Media and
 470 // system messages remain text, including the written sender.
 471 func whatsappImportRead(path, owner string, location *time.Location) (chat WhatsAppChat, messages []WhatsAppMessage, err error) {
 472         chat.Name = strings.TrimPrefix(strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)), "WhatsApp Chat - ")
 473         if chat.Name == "" || strings.ContainsAny(chat.Name, "\r\n") {
 474                 return chat, nil, errors.New("read export: chat name must be one nonempty line")
 475         }
 476
 477         archive, err := zip.OpenReader(path)
 478         if err != nil {
 479                 return chat, nil, fmt.Errorf("open export: %w", err)
 480         }
 481
 482         defer func() {
 483                 closeErr := archive.Close()
 484                 if closeErr != nil {
 485                         err = errors.Join(err, fmt.Errorf("close export: %w", closeErr))
 486                 }
 487         }()
 488
 489         var file *zip.File
 490
 491         for _, entry := range archive.File {
 492                 if filepath.Base(entry.Name) == "_chat.txt" && !entry.FileInfo().IsDir() {
 493                         if file != nil {
 494                                 return chat, nil, errors.New("read export: more than one _chat.txt")
 495                         }
 496
 497                         file = entry
 498                 }
 499         }
 500
 501         if file == nil {
 502                 return chat, nil, errors.New("read export: _chat.txt is missing")
 503         }
 504
 505         if file.UncompressedSize64 > whatsappImportBytes {
 506                 return chat, nil, errors.New("read export: chat text exceeds 64 MiB")
 507         }
 508
 509         reader, err := file.Open()
 510         if err != nil {
 511                 return chat, nil, fmt.Errorf("open chat text: %w", err)
 512         }
 513
 514         defer func() {
 515                 closeErr := reader.Close()
 516                 if closeErr != nil {
 517                         err = errors.Join(err, fmt.Errorf("close export chat text: %w", closeErr))
 518                 }
 519         }()
 520
 521         scanner := bufio.NewScanner(io.LimitReader(reader, whatsappImportBytes+1))
 522         scanner.Buffer(make([]byte, 64*1024), buildLineBytes)
 523
 524         for scanner.Scan() {
 525                 line := strings.ReplaceAll(scanner.Text(), "\u200e", "")
 526                 match := whatsappImportStart.FindStringSubmatch(line)
 527                 if match == nil {
 528                         if len(messages) == 0 {
 529                                 if line == "" {
 530                                         continue
 531                                 }
 532
 533                                 return chat, nil, errors.New("read export: text before the first dated message")
 534                         }
 535
 536                         messages[len(messages)-1].Text += "\n" + line
 537
 538                         continue
 539                 }
 540
 541                 var date time.Time
 542
 543                 date, err = time.ParseInLocation("02/01/2006, 15:04:05", match[1], location)
 544                 if err != nil {
 545                         return chat, nil, fmt.Errorf("parse export local date: %w", err)
 546                 }
 547
 548                 messages = append(messages, WhatsAppMessage{
 549                         Sender: match[2],
 550                         Text:   match[3],
 551                         Time:   date,
 552                         Mine:   owner != "" && match[2] == owner && match[2] != chat.Name,
 553                 })
 554
 555                 if match[2] == chat.Name {
 556                         chat.Group = true
 557                 }
 558         }
 559
 560         if scanner.Err() != nil {
 561                 return chat, nil, fmt.Errorf("read chat text: %w", scanner.Err())
 562         }
 563
 564         if len(messages) == 0 {
 565                 return chat, nil, errors.New("read export: no dated messages")
 566         }
 567
 568         slices.SortStableFunc(messages, func(a, b WhatsAppMessage) int {
 569                 return a.Time.Compare(b.Time)
 570         })
 571
 572         return chat, messages, nil
 573 }