Refactor/cleanup code (whatsapp)

3 files changed, 97 insertions, 101 deletions

diff --git a/bridge/whatsapp/handlers.go b/bridge/whatsapp/handlers.go
index 4d1c9a13..8b25c146 100644
--- a/bridge/whatsapp/handlers.go
+++ b/bridge/whatsapp/handlers.go
@@ -24,7 +24,8 @@ Check:
 func (b *Bwhatsapp) HandleError(err error) {
 	// ignore received invalid data errors. https://github.com/42wim/matterbridge/issues/843
 	// ignore tag 174 errors. https://github.com/42wim/matterbridge/issues/1094
-	if strings.Contains(err.Error(), "error processing data: received invalid data") || strings.Contains(err.Error(), "invalid string with tag 174") {
+	if strings.Contains(err.Error(), "error processing data: received invalid data") ||
+		strings.Contains(err.Error(), "invalid string with tag 174") {
 		return
 	}
 
@@ -47,16 +48,22 @@ func (b *Bwhatsapp) reconnect(err error) {
 		Max:    5 * time.Minute,
 		Jitter: true,
 	}
+
 	for {
 		d := bf.Duration()
+
 		b.Log.Errorf("Connection failed, underlying error: %v", err)
 		b.Log.Infof("Waiting %s...", d)
+
 		time.Sleep(d)
+
 		b.Log.Info("Reconnecting...")
+
 		err := b.conn.Restore()
 		if err == nil {
 			bf.Reset()
 			b.startedAt = uint64(time.Now().Unix())
+
 			return
 		}
 	}
@@ -64,7 +71,7 @@ func (b *Bwhatsapp) reconnect(err error) {
 
 // HandleTextMessage sent from WhatsApp, relay it to the brige
 func (b *Bwhatsapp) HandleTextMessage(message whatsapp.TextMessage) {
-	if message.Info.FromMe { // || !strings.Contains(strings.ToLower(message.Text), "@echo") {
+	if message.Info.FromMe {
 		return
 	}
 	// whatsapp sends last messages to show context , cut them
@@ -101,11 +108,13 @@ func (b *Bwhatsapp) HandleTextMessage(message whatsapp.TextMessage) {
 			if mention == "" {
 				mention = "someone"
 			}
+
 			message.Text = strings.Replace(message.Text, "@"+numberAndSuffix[0], "@"+mention, 1)
 		}
 	}
 
 	b.Log.Debugf("<= Sending message from %s on %s to gateway", senderJID, b.Account)
+
 	rmsg := config.Message{
 		UserID:    senderJID,
 		Username:  senderName,
@@ -116,15 +125,15 @@ func (b *Bwhatsapp) HandleTextMessage(message whatsapp.TextMessage) {
 		Protocol:  b.Protocol,
 		Extra:     make(map[string][]interface{}),
 		//	ParentID: TODO, // TODO handle thread replies  // map from Info.QuotedMessageID string
-		//	Event     string    `json:"event"`
-		//	Gateway   string  // will be added during message processing
-		ID: message.Info.Id}
+		ID: message.Info.Id,
+	}
 
 	if avatarURL, exists := b.userAvatars[senderJID]; exists {
 		rmsg.Avatar = avatarURL
 	}
 
 	b.Log.Debugf("<= Message is %#v", rmsg)
+
 	b.Remote <- rmsg
 }
 
@@ -157,6 +166,7 @@ func (b *Bwhatsapp) HandleImageMessage(message whatsapp.ImageMessage) {
 	}
 
 	b.Log.Debugf("<= Sending message from %s on %s to gateway", senderJID, b.Account)
+
 	rmsg := config.Message{
 		UserID:    senderJID,
 		Username:  senderName,
@@ -166,9 +176,8 @@ func (b *Bwhatsapp) HandleImageMessage(message whatsapp.ImageMessage) {
 		Protocol:  b.Protocol,
 		Extra:     make(map[string][]interface{}),
 		//  ParentID: TODO,      // TODO handle thread replies  // map from Info.QuotedMessageID string
-		//  Event     string    `json:"event"`
-		//  Gateway   string     // will be added during message processing
-		ID: message.Info.Id}
+		ID: message.Info.Id,
+	}
 
 	if avatarURL, exists := b.userAvatars[senderJID]; exists {
 		rmsg.Avatar = avatarURL
@@ -196,6 +205,7 @@ func (b *Bwhatsapp) HandleImageMessage(message whatsapp.ImageMessage) {
 	helper.HandleDownloadData(b.Log, &rmsg, filename, message.Caption, "", &data, b.General)
 
 	b.Log.Debugf("<= Image Message is %#v", rmsg)
+
 	b.Remote <- rmsg
 }
 
diff --git a/bridge/whatsapp/helpers.go b/bridge/whatsapp/helpers.go
index aa5f3188..e424387d 100644
--- a/bridge/whatsapp/helpers.go
+++ b/bridge/whatsapp/helpers.go
@@ -6,22 +6,24 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"strings"
 
 	qrcodeTerminal "github.com/Baozisoftware/qrcode-terminal-go"
 	"github.com/Rhymen/go-whatsapp"
 )
 
 type ProfilePicInfo struct {
-	URL string `json:"eurl"`
-	Tag string `json:"tag"`
-
-	Status int16 `json:"status"`
+	URL    string `json:"eurl"`
+	Tag    string `json:"tag"`
+	Status int16  `json:"status"`
 }
 
 func qrFromTerminal(invert bool) chan string {
 	qr := make(chan string)
+
 	go func() {
 		terminal := qrcodeTerminal.New()
+
 		if invert {
 			terminal = qrcodeTerminal.New2(qrcodeTerminal.ConsoleColors.BrightWhite, qrcodeTerminal.ConsoleColors.BrightBlack, qrcodeTerminal.QRCodeRecoveryLevels.Medium)
 		}
@@ -44,13 +46,12 @@ func (b *Bwhatsapp) readSession() (whatsapp.Session, error) {
 	if err != nil {
 		return session, err
 	}
+
 	defer file.Close()
+
 	decoder := gob.NewDecoder(file)
-	err = decoder.Decode(&session)
-	if err != nil {
-		return session, err
-	}
-	return session, nil
+
+	return session, decoder.Decode(&session)
 }
 
 func (b *Bwhatsapp) writeSession(session whatsapp.Session) error {
@@ -65,11 +66,31 @@ func (b *Bwhatsapp) writeSession(session whatsapp.Session) error {
 	if err != nil {
 		return err
 	}
+
 	defer file.Close()
+
 	encoder := gob.NewEncoder(file)
-	err = encoder.Encode(session)
 
-	return err
+	return encoder.Encode(session)
+}
+
+func (b *Bwhatsapp) restoreSession() (*whatsapp.Session, error) {
+	session, err := b.readSession()
+	if err != nil {
+		b.Log.Warn(err.Error())
+	}
+
+	b.Log.Debugln("Restoring WhatsApp session..")
+
+	session, err = b.conn.RestoreWithSession(session)
+	if err != nil {
+		// restore session connection timed out (I couldn't get over it without logging in again)
+		return nil, errors.New("failed to restore session: " + err.Error())
+	}
+
+	b.Log.Debugln("Session restored successfully!")
+
+	return &session, nil
 }
 
 func (b *Bwhatsapp) getSenderName(senderJid string) string {
@@ -114,6 +135,7 @@ func (b *Bwhatsapp) getSenderNotify(senderJid string) string {
 	if sender, exists := b.users[senderJid]; exists {
 		return sender.Notify
 	}
+
 	return ""
 }
 
@@ -122,11 +144,20 @@ func (b *Bwhatsapp) GetProfilePicThumb(jid string) (*ProfilePicInfo, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to get avatar: %v", err)
 	}
+
 	content := <-data
 	info := &ProfilePicInfo{}
+
 	err = json.Unmarshal([]byte(content), info)
 	if err != nil {
 		return info, fmt.Errorf("failed to unmarshal avatar info: %v", err)
 	}
+
 	return info, nil
 }
+
+func isGroupJid(identifier string) bool {
+	return strings.HasSuffix(identifier, "@g.us") ||
+		strings.HasSuffix(identifier, "@temp") ||
+		strings.HasSuffix(identifier, "@broadcast")
+}
diff --git a/bridge/whatsapp/whatsapp.go b/bridge/whatsapp/whatsapp.go
index ae634d6a..8c2c9cfe 100644
--- a/bridge/whatsapp/whatsapp.go
+++ b/bridge/whatsapp/whatsapp.go
@@ -28,7 +28,6 @@ const (
 type Bwhatsapp struct {
 	*bridge.Config
 
-	// https://github.com/Rhymen/go-whatsapp/blob/c31092027237441cffba1b9cb148eadf7c83c3d2/session.go#L18-L21
 	session   *whatsapp.Session
 	conn      *whatsapp.Conn
 	startedAt uint64
@@ -40,6 +39,7 @@ type Bwhatsapp struct {
 // New Create a new WhatsApp bridge. This will be called for each [whatsapp.<server>] entry you have in the config file
 func New(cfg *bridge.Config) bridge.Bridger {
 	number := cfg.GetString(cfgNumber)
+
 	if number == "" {
 		cfg.Log.Fatalf("Missing configuration for WhatsApp bridge: Number")
 	}
@@ -50,21 +50,17 @@ func New(cfg *bridge.Config) bridge.Bridger {
 		users:       make(map[string]whatsapp.Contact),
 		userAvatars: make(map[string]string),
 	}
+
 	return b
 }
 
 // Connect to WhatsApp. Required implementation of the Bridger interface
-// https://github.com/42wim/matterbridge/blob/2cfd880cdb0df29771bf8f31df8d990ab897889d/bridge/bridge.go#L11-L16
 func (b *Bwhatsapp) Connect() error {
-	b.RLock() // TODO do we need locking for Whatsapp?
-	defer b.RUnlock()
-
 	number := b.GetString(cfgNumber)
 	if number == "" {
-		return errors.New("WhatsApp's telephone Number need to be configured")
+		return errors.New("whatsapp's telephone number need to be configured")
 	}
 
-	// https://github.com/Rhymen/go-whatsapp#creating-a-connection
 	b.Log.Debugln("Connecting to WhatsApp..")
 	conn, err := whatsapp.NewConn(20 * time.Second)
 	if err != nil {
@@ -77,35 +73,18 @@ func (b *Bwhatsapp) Connect() error {
 	b.Log.Debugln("WhatsApp connection successful")
 
 	// load existing session in order to keep it between restarts
-	if b.session == nil {
-		var session whatsapp.Session
-		session, err = b.readSession()
-
-		if err == nil {
-			b.Log.Debugln("Restoring WhatsApp session..")
-
-			// https://github.com/Rhymen/go-whatsapp#restore
-			session, err = b.conn.RestoreWithSession(session)
-			if err != nil {
-				// TODO return or continue to normal login?
-				// restore session connection timed out (I couldn't get over it without logging in again)
-				return errors.New("failed to restore session: " + err.Error())
-			}
-
-			b.session = &session
-			b.Log.Debugln("Session restored successfully!")
-		} else {
-			b.Log.Warn(err.Error())
-		}
+	b.session, err = b.restoreSession()
+	if err != nil {
+		b.Log.Warn(err.Error())
 	}
 
 	// login to a new session
 	if b.session == nil {
-		err = b.Login()
-		if err != nil {
+		if err = b.Login(); err != nil {
 			return err
 		}
 	}
+
 	b.startedAt = uint64(time.Now().Unix())
 
 	_, err = b.conn.Contacts()
@@ -116,6 +95,7 @@ func (b *Bwhatsapp) Connect() error {
 	// see https://github.com/Rhymen/go-whatsapp/issues/137#issuecomment-480316013
 	for len(b.conn.Store.Contacts) == 0 {
 		b.conn.Contacts() // nolint:errcheck
+
 		<-time.After(1 * time.Second)
 	}
 
@@ -135,12 +115,13 @@ func (b *Bwhatsapp) Connect() error {
 			info, err := b.GetProfilePicThumb(jid)
 			if err != nil {
 				b.Log.Warnf("Could not get profile photo of %s: %v", jid, err)
-
 			} else {
-				// TODO any race conditions here?
+				b.Lock()
 				b.userAvatars[jid] = info.URL
+				b.Unlock()
 			}
 		}
+
 		b.Log.Debug("Finished getting avatars..")
 	}()
 
@@ -157,8 +138,10 @@ func (b *Bwhatsapp) Login() error {
 	session, err := b.conn.Login(qrChan)
 	if err != nil {
 		b.Log.Warnln("Failed to log in:", err)
+
 		return err
 	}
+
 	b.session = &session
 
 	b.Log.Infof("Logged into session: %#v", session)
@@ -169,29 +152,17 @@ func (b *Bwhatsapp) Login() error {
 		fmt.Fprintf(os.Stderr, "error saving session: %v\n", err)
 	}
 
-	// TODO change connection strings to configured ones longClientName:"github.com/rhymen/go-whatsapp", shortClientName:"go-whatsapp"}" prefix=whatsapp
-	// TODO get also a nice logo
-
-	// TODO notification about unplugged and dead battery
-	// conn.Info: Wid, Pushname, Connected, Battery, Plugged
-
 	return nil
 }
 
 // Disconnect is called while reconnecting to the bridge
-// TODO 42wim Documentation would be helpful on when reconnects happen and what should be done in this function
 // Required implementation of the Bridger interface
-// https://github.com/42wim/matterbridge/blob/2cfd880cdb0df29771bf8f31df8d990ab897889d/bridge/bridge.go#L11-L16
 func (b *Bwhatsapp) Disconnect() error {
 	// We could Logout, but that would close the session completely and would require a new QR code scan
 	// https://github.com/Rhymen/go-whatsapp/blob/c31092027237441cffba1b9cb148eadf7c83c3d2/session.go#L377-L381
 	return nil
 }
 
-func isGroupJid(identifier string) bool {
-	return strings.HasSuffix(identifier, "@g.us") || strings.HasSuffix(identifier, "@temp") || strings.HasSuffix(identifier, "@broadcast")
-}
-
 // JoinChannel Join a WhatsApp group specified in gateway config as channel='number-id@g.us' or channel='Channel name'
 // Required implementation of the Bridger interface
 // https://github.com/42wim/matterbridge/blob/2cfd880cdb0df29771bf8f31df8d990ab897889d/bridge/bridge.go#L11-L16
@@ -210,39 +181,33 @@ func (b *Bwhatsapp) JoinChannel(channel config.ChannelInfo) error {
 		if _, exists := b.conn.Store.Contacts[channel.Name]; !exists {
 			return fmt.Errorf("account doesn't belong to group with jid %s", channel.Name)
 		}
-	} else {
-		// channel.Name specifies group name that might change, warn about it
-		var jids []string
-		for id, contact := range b.conn.Store.Contacts {
-			if isGroupJid(id) && contact.Name == channel.Name {
-				jids = append(jids, id)
-			}
-		}
 
-		switch len(jids) {
-		case 0:
-			// didn't match any group - print out possibilites
-			// TODO sort
-			// copy b;
-			//sort.Slice(people, func(i, j int) bool {
-			//	return people[i].Age > people[j].Age
-			//})
-			for id, contact := range b.conn.Store.Contacts {
-				if isGroupJid(id) {
-					b.Log.Infof("%s %s", contact.Jid, contact.Name)
-				}
-			}
-			return fmt.Errorf("please specify group's JID from the list above instead of the name '%s'", channel.Name)
-
-		case 1:
-			return fmt.Errorf("group name might change. Please configure gateway with channel=\"%v\" instead of channel=\"%v\"", jids[0], channel.Name)
+		return nil
+	}
 
-		default:
-			return fmt.Errorf("there is more than one group with name '%s'. Please specify one of JIDs as channel name: %v", channel.Name, jids)
+	// channel.Name specifies group name that might change, warn about it
+	var jids []string
+	for id, contact := range b.conn.Store.Contacts {
+		if isGroupJid(id) && contact.Name == channel.Name {
+			jids = append(jids, id)
 		}
 	}
 
-	return nil
+	switch len(jids) {
+	case 0:
+		// didn't match any group - print out possibilites
+		for id, contact := range b.conn.Store.Contacts {
+			if isGroupJid(id) {
+				b.Log.Infof("%s %s", contact.Jid, contact.Name)
+			}
+		}
+
+		return fmt.Errorf("please specify group's JID from the list above instead of the name '%s'", channel.Name)
+	case 1:
+		return fmt.Errorf("group name might change. Please configure gateway with channel=\"%v\" instead of channel=\"%v\"", jids[0], channel.Name)
+	default:
+		return fmt.Errorf("there is more than one group with name '%s'. Please specify one of JIDs as channel name: %v", channel.Name, jids)
+	}
 }
 
 // Post a document message from the bridge to WhatsApp
@@ -331,7 +296,6 @@ func (b *Bwhatsapp) Send(msg config.Message) (string, error) {
 		b.Log.Debugf("updating message with id %s", msg.ID)
 
 		msg.Text += " (edited)"
-		// TODO handle edit as a message reply with updated text
 	}
 
 	// Handle Upload a file
@@ -361,16 +325,7 @@ func (b *Bwhatsapp) Send(msg config.Message) (string, error) {
 
 	b.Log.Debugf("=> Sending %#v", msg)
 
-	// create message ID
-	// TODO follow and act if https://github.com/Rhymen/go-whatsapp/issues/101 implemented
-	idBytes := make([]byte, 10)
-	if _, err := rand.Read(idBytes); err != nil {
-		b.Log.Warn(err.Error())
-	}
-	message.Info.Id = strings.ToUpper(hex.EncodeToString(idBytes))
-	_, err := b.conn.Send(message)
-
-	return message.Info.Id, err
+	return b.conn.Send(message)
 }
 
 // TODO do we want that? to allow login with QR code from a bridged channel? https://github.com/tulir/mautrix-whatsapp/blob/513eb18e2d59bada0dd515ee1abaaf38a3bfe3d5/commands.go#L76