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, buildCallLimit),
 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         err = json.Unmarshal(data, &state)
 300         if err != nil {
 301                 return nil, "", fmt.Errorf("decode pending import names: %w", err)
 302         }
 303
 304         for _, batch := range state.Batches {
 305                 whatsappImportName(names, batch.Chat)
 306         }
 307
 308         return names, owner, nil
 309 }
 310
 311 // whatsappImportName adds an exact match unless distinct IDs share its name.
 312 func whatsappImportName(names map[string]WhatsAppChat, chat WhatsAppChat) {
 313         old, found := names[chat.Name]
 314         if found && old.ID != chat.ID {
 315                 chat.ID = ""
 316         }
 317
 318         names[chat.Name] = chat
 319 }
 320
 321 // whatsappImportPlan filters exports at the first message of the earliest live
 322 // batch and skips chat-days already imported. Each remaining chat-day becomes a
 323 // work message dated at its last exported message. Overlapping chat-days in the
 324 // inputs are refused, not silently combined or duplicated.
 325 func whatsappImportPlan(ctx context.Context, paths []string, old []DatedMessage, names map[string]WhatsAppChat, owner string, location *time.Location) (WhatsAppImportPlan, error) {
 326         plan := WhatsAppImportPlan{ownerKnown: owner != ""}
 327         names = maps.Clone(names)
 328         if names == nil {
 329                 names = make(map[string]WhatsAppChat)
 330         }
 331
 332         imported := make(map[string]bool)
 333         for _, entry := range old {
 334                 if entry.Kind != "work" || !strings.HasPrefix(entry.Text, "[WhatsApp] ") {
 335                         continue
 336                 }
 337
 338                 first, _, _ := strings.Cut(entry.Text, "\n")
 339                 marked := strings.HasSuffix(first, whatsappImportMark)
 340                 first = strings.TrimSuffix(first, whatsappImportMark)
 341                 match := whatsappImportHeader.FindStringSubmatch(first)
 342                 name, day := "", ""
 343                 if len(match) > 0 {
 344                         name, day = match[1], match[3]
 345                         whatsappImportName(names, WhatsAppChat{
 346                                 ID:    match[2],
 347                                 Name:  name,
 348                                 Group: strings.Contains(first, ", group "+match[2]+","),
 349                         })
 350                 } else if marked {
 351                         match = whatsappImportUnknown.FindStringSubmatch(first)
 352                         if len(match) > 0 {
 353                                 name, day = match[1], match[2]
 354                         }
 355                 }
 356
 357                 if name == "" {
 358                         return plan, errors.New("read WhatsApp history: invalid batch header")
 359                 }
 360
 361                 if marked {
 362                         imported[name+"\n"+day] = true
 363                         continue
 364                 }
 365
 366                 _, body, found := strings.Cut(entry.Text, "\n")
 367                 if !found || len(body) < len("15:04 ") {
 368                         return plan, errors.New("read WhatsApp history: missing first message time")
 369                 }
 370
 371                 date, err := time.ParseInLocation("2006-01-02 15:04", day+" "+body[:5], location)
 372                 if err != nil {
 373                         return plan, fmt.Errorf("parse first live message time: %w", err)
 374                 }
 375
 376                 if plan.cutoff.IsZero() || date.Before(plan.cutoff) {
 377                         plan.cutoff = date
 378                 }
 379         }
 380
 381         seen := make(map[string]bool)
 382         for _, path := range paths {
 383                 if ctx.Err() != nil {
 384                         return plan, ctx.Err()
 385                 }
 386
 387                 chat, messages, err := whatsappImportRead(path, owner, location)
 388                 if err != nil {
 389                         return plan, fmt.Errorf("read import export: %w", err)
 390                 }
 391
 392                 known, found := names[chat.Name]
 393                 if found {
 394                         chat = known
 395                 }
 396
 397                 if !slices.Contains(plan.chats, chat.Name) {
 398                         plan.chats = append(plan.chats, chat.Name)
 399                 }
 400
 401                 days := make(map[string][]WhatsAppMessage)
 402                 for _, message := range messages {
 403                         if !plan.cutoff.IsZero() && !message.Time.Before(plan.cutoff) {
 404                                 continue
 405                         }
 406
 407                         day := message.Time.In(location).Format(time.DateOnly)
 408                         days[day] = append(days[day], message)
 409                 }
 410
 411                 for _, day := range slices.Sorted(maps.Keys(days)) {
 412                         key := chat.Name + "\n" + day
 413                         if imported[key] {
 414                                 plan.skipped++
 415                                 continue
 416                         }
 417
 418                         if seen[key] {
 419                                 return plan, errors.New("import inputs overlap: give only one export for each chat and day")
 420                         }
 421
 422                         seen[key] = true
 423                         batch := WhatsAppBatch{
 424                                 Chat:     chat,
 425                                 Messages: days[day],
 426                         }
 427
 428                         text := batch.Text(location)
 429                         first, body, _ := strings.Cut(text, "\n")
 430                         if chat.ID == "" {
 431                                 first = "[WhatsApp] " + chat.Name + ", " + batch.Messages[0].Time.In(location).Format(whatsappBatchDayLayout)
 432                         }
 433
 434                         text = first + whatsappImportMark + "\n" + body
 435                         plan.added = append(plan.added, DatedMessage{
 436                                 Kind: "work",
 437                                 Text: text,
 438                                 Date: batch.Messages[len(batch.Messages)-1].Time.Format(time.RFC3339Nano),
 439                         })
 440
 441                         plan.bytes += len(text)
 442                         firstDate, lastDate := batch.Messages[0].Time, batch.Messages[len(batch.Messages)-1].Time
 443                         if plan.first.IsZero() || firstDate.Before(plan.first) {
 444                                 plan.first = firstDate
 445                         }
 446
 447                         if plan.last.IsZero() || lastDate.After(plan.last) {
 448                                 plan.last = lastDate
 449                         }
 450                 }
 451         }
 452
 453         var err error
 454         plan.added, err = rebuildMerge(nil, plan.added)
 455         if err != nil {
 456                 return plan, fmt.Errorf("order import batches: %w", err)
 457         }
 458
 459         plan.messages = len(old) + len(plan.added)
 460         slices.Sort(plan.chats)
 461
 462         return plan, nil
 463 }
 464
 465 // whatsappImportRead reads the single _chat.txt in an export without extracting
 466 // archive paths. Non-start lines continue the preceding message, and all
 467 // invisible left-to-right marks are removed. Dates use location. Media and
 468 // system messages remain text, including the written sender.
 469 func whatsappImportRead(path, owner string, location *time.Location) (chat WhatsAppChat, messages []WhatsAppMessage, err error) {
 470         chat.Name = strings.TrimPrefix(strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)), "WhatsApp Chat - ")
 471         if chat.Name == "" || strings.ContainsAny(chat.Name, "\r\n") {
 472                 return chat, nil, errors.New("read export: chat name must be one nonempty line")
 473         }
 474
 475         archive, err := zip.OpenReader(path)
 476         if err != nil {
 477                 return chat, nil, fmt.Errorf("open export: %w", err)
 478         }
 479
 480         defer func() {
 481                 closeErr := archive.Close()
 482                 if closeErr != nil {
 483                         err = errors.Join(err, fmt.Errorf("close export: %w", closeErr))
 484                 }
 485         }()
 486
 487         var file *zip.File
 488         for _, entry := range archive.File {
 489                 if filepath.Base(entry.Name) == "_chat.txt" && !entry.FileInfo().IsDir() {
 490                         if file != nil {
 491                                 return chat, nil, errors.New("read export: more than one _chat.txt")
 492                         }
 493
 494                         file = entry
 495                 }
 496         }
 497
 498         if file == nil {
 499                 return chat, nil, errors.New("read export: _chat.txt is missing")
 500         }
 501
 502         if file.UncompressedSize64 > whatsappImportBytes {
 503                 return chat, nil, errors.New("read export: chat text exceeds 64 MiB")
 504         }
 505
 506         reader, err := file.Open()
 507         if err != nil {
 508                 return chat, nil, fmt.Errorf("open chat text: %w", err)
 509         }
 510
 511         defer func() {
 512                 closeErr := reader.Close()
 513                 if closeErr != nil {
 514                         err = errors.Join(err, fmt.Errorf("close export chat text: %w", closeErr))
 515                 }
 516         }()
 517
 518         scanner := bufio.NewScanner(io.LimitReader(reader, whatsappImportBytes+1))
 519         scanner.Buffer(make([]byte, 64*1024), buildLineBytes)
 520         for scanner.Scan() {
 521                 line := strings.ReplaceAll(scanner.Text(), "\u200e", "")
 522                 match := whatsappImportStart.FindStringSubmatch(line)
 523                 if match == nil {
 524                         if len(messages) == 0 {
 525                                 if line == "" {
 526                                         continue
 527                                 }
 528
 529                                 return chat, nil, errors.New("read export: text before the first dated message")
 530                         }
 531
 532                         messages[len(messages)-1].Text += "\n" + line
 533
 534                         continue
 535                 }
 536
 537                 var date time.Time
 538
 539                 date, err = time.ParseInLocation("02/01/2006, 15:04:05", match[1], location)
 540                 if err != nil {
 541                         return chat, nil, fmt.Errorf("parse export local date: %w", err)
 542                 }
 543
 544                 messages = append(messages, WhatsAppMessage{
 545                         Sender: match[2],
 546                         Text:   match[3],
 547                         Time:   date,
 548                         Mine:   owner != "" && match[2] == owner && match[2] != chat.Name,
 549                 })
 550
 551                 if match[2] == chat.Name {
 552                         chat.Group = true
 553                 }
 554         }
 555
 556         if scanner.Err() != nil {
 557                 return chat, nil, fmt.Errorf("read chat text: %w", scanner.Err())
 558         }
 559
 560         if len(messages) == 0 {
 561                 return chat, nil, errors.New("read export: no dated messages")
 562         }
 563
 564         slices.SortStableFunc(messages, func(a, b WhatsAppMessage) int {
 565                 return a.Time.Compare(b.Time)
 566         })
 567
 568         return chat, messages, nil
 569 }