From 8fc9fc7bc0ce27278defb4cef8a3dec208b86207 Mon Sep 17 00:00:00 2001 From: Aldino Kemal Date: Sat, 1 Mar 2025 07:02:36 +0700 Subject: [PATCH] refactor: Improve WhatsApp event handling and code organization - Refactor `init.go` to break down complex event handler into smaller, focused functions - Add type definitions and comments to improve code readability - Extract event handling logic into separate functions for better maintainability - Improve error handling and logging in database initialization - Add `FormatJID` utility function to handle JID formatting - Update README with minor formatting improvements --- readme.md | 4 +- src/pkg/whatsapp/init.go | 323 ++++++++++++++++++++++++-------------- src/pkg/whatsapp/utils.go | 12 ++ 3 files changed, 223 insertions(+), 116 deletions(-) diff --git a/readme.md b/readme.md index 67267bf..87f1551 100644 --- a/readme.md +++ b/readme.md @@ -38,7 +38,8 @@ Now that we support ARM64 for Linux: - `--webhook="http://yourwebhook.site/handler"`, or you can simplify - `-w="http://yourwebhook.site/handler"` - Webhook Secret - Our webhook will be sent to you with an HMAC header and a sha256 default key `secret`.
+ Our webhook will be sent to you with an HMAC header and a sha256 default key `secret`. + You may modify this by using the option below: - `--webhook-secret="secret"` @@ -53,6 +54,7 @@ You can configure the application using either command-line flags (shown above) ### Environment Variables To use environment variables: + 1. Copy `.env.example` to `.env` in your project root 2. Modify the values in `.env` according to your needs 3. Or set the same variables as system environment variables diff --git a/src/pkg/whatsapp/init.go b/src/pkg/whatsapp/init.go index c81835a..f5ecc4f 100644 --- a/src/pkg/whatsapp/init.go +++ b/src/pkg/whatsapp/init.go @@ -25,13 +25,7 @@ import ( "google.golang.org/protobuf/proto" ) -var ( - cli *whatsmeow.Client - log waLog.Logger - historySyncID int32 - startupTime = time.Now().Unix() -) - +// Type definitions type ExtractedMedia struct { MediaPath string `json:"media_path"` MimeType string `json:"mime_type"` @@ -50,29 +44,40 @@ type evtMessage struct { QuotedMessage string `json:"quoted_message,omitempty"` } +// Global variables +var ( + cli *whatsmeow.Client + log waLog.Logger + historySyncID int32 + startupTime = time.Now().Unix() +) + +// InitWaDB initializes the WhatsApp database connection func InitWaDB() *sqlstore.Container { - // Running Whatsapp log = waLog.Stdout("Main", config.WhatsappLogLevel, true) dbLog := waLog.Stdout("Database", config.WhatsappLogLevel, true) - var storeContainer *sqlstore.Container - var err error + storeContainer, err := initDatabase(dbLog) + if err != nil { + log.Errorf("Database initialization error: %v", err) + panic(pkgError.InternalServerError(fmt.Sprintf("Database initialization error: %v", err))) + } + + return storeContainer +} + +// initDatabase creates and returns a database store container based on the configured URI +func initDatabase(dbLog waLog.Logger) (*sqlstore.Container, error) { if strings.HasPrefix(config.DBURI, "file:") { - storeContainer, err = sqlstore.New("sqlite3", config.DBURI, dbLog) + return sqlstore.New("sqlite3", config.DBURI, dbLog) } else if strings.HasPrefix(config.DBURI, "postgres:") { - storeContainer, err = sqlstore.New("postgres", config.DBURI, dbLog) - } else { - log.Errorf("Unknown database type: %s", config.DBURI) - panic(pkgError.InternalServerError(fmt.Sprintf("Unknown database type: %s. Currently only sqlite3(file:) and postgres are supported", config.DBURI))) + return sqlstore.New("postgres", config.DBURI, dbLog) } - if err != nil { - log.Errorf("Failed to connect to database: %v", err) - panic(pkgError.InternalServerError(fmt.Sprintf("Failed to connect to database: %v", err))) - } - return storeContainer + return nil, fmt.Errorf("unknown database type: %s. Currently only sqlite3(file:) and postgres are supported", config.DBURI) } +// InitWaCLI initializes the WhatsApp client func InitWaCLI(storeContainer *sqlstore.Container) *whatsmeow.Client { device, err := storeContainer.GetFirstDevice() if err != nil { @@ -85,9 +90,12 @@ func InitWaCLI(storeContainer *sqlstore.Container) *whatsmeow.Client { panic("No device found") } + // Configure device properties osName := fmt.Sprintf("%s %s", config.AppOs, config.AppVersion) store.DeviceProps.PlatformType = &config.AppPlatform store.DeviceProps.Os = &osName + + // Create and configure the client cli = whatsmeow.NewClient(device, waLog.Stdout("Client", config.WhatsappLogLevel, true)) cli.EnableAutoReconnect = true cli.AutoTrustIdentity = true @@ -96,120 +104,205 @@ func InitWaCLI(storeContainer *sqlstore.Container) *whatsmeow.Client { return cli } +// handler is the main event handler for WhatsApp events func handler(rawEvt interface{}) { switch evt := rawEvt.(type) { case *events.DeleteForMe: - log.Infof("Deleted message %s for %s", evt.MessageID, evt.SenderJID.String()) + handleDeleteForMe(evt) case *events.AppStateSyncComplete: - if len(cli.Store.PushName) > 0 && evt.Name == appstate.WAPatchCriticalBlock { - err := cli.SendPresence(types.PresenceAvailable) - if err != nil { - log.Warnf("Failed to send available presence: %v", err) - } else { - log.Infof("Marked self as available") - } - } + handleAppStateSyncComplete(evt) case *events.PairSuccess: - websocket.Broadcast <- websocket.BroadcastMessage{ - Code: "LOGIN_SUCCESS", - Message: fmt.Sprintf("Successfully pair with %s", evt.ID.String()), - } + handlePairSuccess(evt) case *events.LoggedOut: - websocket.Broadcast <- websocket.BroadcastMessage{ - Code: "LIST_DEVICES", - Result: nil, - } + handleLoggedOut() case *events.Connected, *events.PushNameSetting: - if len(cli.Store.PushName) == 0 { - return - } + handleConnectionEvents() + case *events.StreamReplaced: + handleStreamReplaced() + case *events.Message: + handleMessage(evt) + case *events.Receipt: + handleReceipt(evt) + case *events.Presence: + handlePresence(evt) + case *events.HistorySync: + handleHistorySync(evt) + case *events.AppState: + handleAppState(evt) + } +} + +// Event handler functions + +func handleDeleteForMe(evt *events.DeleteForMe) { + log.Infof("Deleted message %s for %s", evt.MessageID, evt.SenderJID.String()) +} - // Send presence available when connecting and when the pushname is changed. - // This makes sure that outgoing messages always have the right pushname. - err := cli.SendPresence(types.PresenceAvailable) - if err != nil { +func handleAppStateSyncComplete(evt *events.AppStateSyncComplete) { + if len(cli.Store.PushName) > 0 && evt.Name == appstate.WAPatchCriticalBlock { + if err := cli.SendPresence(types.PresenceAvailable); err != nil { log.Warnf("Failed to send available presence: %v", err) } else { log.Infof("Marked self as available") } - case *events.StreamReplaced: - os.Exit(0) - case *events.Message: - metaParts := []string{ - fmt.Sprintf("pushname: %s", evt.Info.PushName), - fmt.Sprintf("timestamp: %s", evt.Info.Timestamp), - } - if evt.Info.Type != "" { - metaParts = append(metaParts, fmt.Sprintf("type: %s", evt.Info.Type)) - } - if evt.Info.Category != "" { - metaParts = append(metaParts, fmt.Sprintf("category: %s", evt.Info.Category)) - } - if evt.IsViewOnce { - metaParts = append(metaParts, "view once") - } + } +} - log.Infof("Received message %s from %s (%s): %+v", evt.Info.ID, evt.Info.SourceString(), strings.Join(metaParts, ", "), evt.Message) - message := ExtractMessageText(evt) - utils.RecordMessage(evt.Info.ID, evt.Info.Sender.String(), message) +func handlePairSuccess(evt *events.PairSuccess) { + websocket.Broadcast <- websocket.BroadcastMessage{ + Code: "LOGIN_SUCCESS", + Message: fmt.Sprintf("Successfully pair with %s", evt.ID.String()), + } +} - if img := evt.Message.GetImageMessage(); img != nil { - if path, err := ExtractMedia(config.PathStorages, img); err != nil { - log.Errorf("Failed to download image: %v", err) - } else { - log.Infof("Image downloaded to %s", path) - } - } +func handleLoggedOut() { + websocket.Broadcast <- websocket.BroadcastMessage{ + Code: "LIST_DEVICES", + Result: nil, + } +} - if config.WhatsappAutoReplyMessage != "" && - !isGroupJid(evt.Info.Chat.String()) && - !strings.Contains(evt.Info.SourceString(), "broadcast") { - _, _ = cli.SendMessage(context.Background(), evt.Info.Sender, &waE2E.Message{Conversation: proto.String(config.WhatsappAutoReplyMessage)}) - } +func handleConnectionEvents() { + if len(cli.Store.PushName) == 0 { + return + } - if len(config.WhatsappWebhook) > 0 && - !strings.Contains(evt.Info.SourceString(), "broadcast") && - !isFromMySelf(evt.Info.SourceString()) { - go func(evt *events.Message) { - if err := forwardToWebhook(evt); err != nil { - logrus.Error("Failed forward to webhook: ", err) - } - }(evt) - } - case *events.Receipt: - if evt.Type == types.ReceiptTypeRead || evt.Type == types.ReceiptTypeReadSelf { - log.Infof("%v was read by %s at %s", evt.MessageIDs, evt.SourceString(), evt.Timestamp) - } else if evt.Type == types.ReceiptTypeDelivered { - log.Infof("%s was delivered to %s at %s", evt.MessageIDs[0], evt.SourceString(), evt.Timestamp) + // Send presence available when connecting and when the pushname is changed. + // This makes sure that outgoing messages always have the right pushname. + if err := cli.SendPresence(types.PresenceAvailable); err != nil { + log.Warnf("Failed to send available presence: %v", err) + } else { + log.Infof("Marked self as available") + } +} + +func handleStreamReplaced() { + os.Exit(0) +} + +func handleMessage(evt *events.Message) { + // Log message metadata + metaParts := buildMessageMetaParts(evt) + log.Infof("Received message %s from %s (%s): %+v", + evt.Info.ID, + evt.Info.SourceString(), + strings.Join(metaParts, ", "), + evt.Message, + ) + + // Record the message + message := ExtractMessageText(evt) + utils.RecordMessage(evt.Info.ID, evt.Info.Sender.String(), message) + + // Handle image message if present + handleImageMessage(evt) + + // Handle auto-reply if configured + handleAutoReply(evt) + + // Forward to webhook if configured + handleWebhookForward(evt) +} + +func buildMessageMetaParts(evt *events.Message) []string { + metaParts := []string{ + fmt.Sprintf("pushname: %s", evt.Info.PushName), + fmt.Sprintf("timestamp: %s", evt.Info.Timestamp), + } + if evt.Info.Type != "" { + metaParts = append(metaParts, fmt.Sprintf("type: %s", evt.Info.Type)) + } + if evt.Info.Category != "" { + metaParts = append(metaParts, fmt.Sprintf("category: %s", evt.Info.Category)) + } + if evt.IsViewOnce { + metaParts = append(metaParts, "view once") + } + return metaParts +} + +func handleImageMessage(evt *events.Message) { + if img := evt.Message.GetImageMessage(); img != nil { + if path, err := ExtractMedia(config.PathStorages, img); err != nil { + log.Errorf("Failed to download image: %v", err) + } else { + log.Infof("Image downloaded to %s", path) } - case *events.Presence: - if evt.Unavailable { - if evt.LastSeen.IsZero() { - log.Infof("%s is now offline", evt.From) - } else { - log.Infof("%s is now offline (last seen: %s)", evt.From, evt.LastSeen) + } +} + +func handleAutoReply(evt *events.Message) { + if config.WhatsappAutoReplyMessage != "" && + !isGroupJid(evt.Info.Chat.String()) && + !evt.Info.IsIncomingBroadcast() && + evt.Message.GetExtendedTextMessage().GetText() != "" { + _, _ = cli.SendMessage( + context.Background(), + FormatJID(evt.Info.Sender.String()), + &waE2E.Message{Conversation: proto.String(config.WhatsappAutoReplyMessage)}, + ) + } +} + +func handleWebhookForward(evt *events.Message) { + if len(config.WhatsappWebhook) > 0 && + !strings.Contains(evt.Info.SourceString(), "broadcast") && + !isFromMySelf(evt.Info.SourceString()) { + go func(evt *events.Message) { + if err := forwardToWebhook(evt); err != nil { + logrus.Error("Failed forward to webhook: ", err) } + }(evt) + } +} + +func handleReceipt(evt *events.Receipt) { + if evt.Type == types.ReceiptTypeRead || evt.Type == types.ReceiptTypeReadSelf { + log.Infof("%v was read by %s at %s", evt.MessageIDs, evt.SourceString(), evt.Timestamp) + } else if evt.Type == types.ReceiptTypeDelivered { + log.Infof("%s was delivered to %s at %s", evt.MessageIDs[0], evt.SourceString(), evt.Timestamp) + } +} + +func handlePresence(evt *events.Presence) { + if evt.Unavailable { + if evt.LastSeen.IsZero() { + log.Infof("%s is now offline", evt.From) } else { - log.Infof("%s is now online", evt.From) - } - case *events.HistorySync: - id := atomic.AddInt32(&historySyncID, 1) - fileName := fmt.Sprintf("%s/history-%d-%s-%d-%s.json", config.PathStorages, startupTime, cli.Store.ID.String(), id, evt.Data.SyncType.String()) - file, err := os.OpenFile(fileName, os.O_WRONLY|os.O_CREATE, 0600) - if err != nil { - log.Errorf("Failed to open file to write history sync: %v", err) - return + log.Infof("%s is now offline (last seen: %s)", evt.From, evt.LastSeen) } - enc := json.NewEncoder(file) - enc.SetIndent("", " ") - err = enc.Encode(evt.Data) - if err != nil { - log.Errorf("Failed to write history sync: %v", err) - return - } - log.Infof("Wrote history sync to %s", fileName) - _ = file.Close() - case *events.AppState: - log.Debugf("App state event: %+v / %+v", evt.Index, evt.SyncActionValue) + } else { + log.Infof("%s is now online", evt.From) } } + +func handleHistorySync(evt *events.HistorySync) { + id := atomic.AddInt32(&historySyncID, 1) + fileName := fmt.Sprintf("%s/history-%d-%s-%d-%s.json", + config.PathStorages, + startupTime, + cli.Store.ID.String(), + id, + evt.Data.SyncType.String(), + ) + + file, err := os.OpenFile(fileName, os.O_WRONLY|os.O_CREATE, 0600) + if err != nil { + log.Errorf("Failed to open file to write history sync: %v", err) + return + } + defer file.Close() + + enc := json.NewEncoder(file) + enc.SetIndent("", " ") + if err = enc.Encode(evt.Data); err != nil { + log.Errorf("Failed to write history sync: %v", err) + return + } + + log.Infof("Wrote history sync to %s", fileName) +} + +func handleAppState(evt *events.AppState) { + log.Debugf("App state event: %+v / %+v", evt.Index, evt.SyncActionValue) +} diff --git a/src/pkg/whatsapp/utils.go b/src/pkg/whatsapp/utils.go index 5f3b9a6..a4e84ca 100644 --- a/src/pkg/whatsapp/utils.go +++ b/src/pkg/whatsapp/utils.go @@ -177,6 +177,18 @@ func MustLogin(waCli *whatsmeow.Client) { } } +func FormatJID(jid string) types.JID { + // Remove any :number suffix if present + if idx := strings.LastIndex(jid, ":"); idx != -1 && strings.Contains(jid, "@s.whatsapp.net") { + jid = jid[:idx] + jid[strings.Index(jid, "@s.whatsapp.net"):] + } + formattedJID, err := ParseJID(jid) + if err != nil { + return types.JID{} + } + return formattedJID +} + // isGroupJid is a helper function to check if the message is from group func isGroupJid(jid string) bool { return strings.Contains(jid, "@g.us")