Browse Source

Use type aliases for different ID types and add puppet type

Tulir Asokan 6 năm trước cách đây
mục cha
commit
edd4f817e4
13 tập tin đã thay đổi với 325 bổ sung189 xóa
  1. 94 3
      config/bridge.go
  2. 2 0
      config/config.go
  3. 1 1
      config/registration.go
  4. 7 6
      database/portal.go
  5. 6 5
      database/puppet.go
  6. 4 3
      database/user.go
  7. 19 3
      example-config.yaml
  8. 23 104
      main.go
  9. 22 58
      matrix.go
  10. 3 2
      portal.go
  11. 113 0
      puppet.go
  12. 26 0
      types/types.go
  13. 5 4
      user.go

+ 94 - 3
config/bridge.go

@@ -19,12 +19,19 @@ package config
 import (
 	"bytes"
 	"text/template"
+	"maunium.net/go/mautrix-appservice"
+	"strings"
+	"strconv"
 )
 
 type BridgeConfig struct {
-	UsernameTemplate    string             `yaml:"username_template"`
-	DisplaynameTemplate string             `yaml:"displayname_template"`
-	StateStore          string             `yaml:"state_store_path"`
+	UsernameTemplate    string `yaml:"username_template"`
+	DisplaynameTemplate string `yaml:"displayname_template"`
+
+	CommandPrefix string `yaml:"command_prefix"`
+
+	Permissions PermissionConfig `yaml:"permissions"`
+
 	usernameTemplate    *template.Template `yaml:"-"`
 	displaynameTemplate *template.Template `yaml:"-"`
 }
@@ -77,3 +84,87 @@ func (bc BridgeConfig) MarshalYAML() (interface{}, error) {
 	bc.UsernameTemplate = bc.FormatUsername("{{.Receiver}}", "{{.UserID}}")
 	return bc, nil
 }
+
+type PermissionConfig map[string]PermissionLevel
+
+type PermissionLevel int
+
+const (
+	PermissionLevelDefault PermissionLevel = 0
+	PermissionLevelUser    PermissionLevel = 10
+	PermissionLevelAdmin   PermissionLevel = 100
+)
+
+func (pc *PermissionConfig) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	rawPC := make(map[string]string)
+	err := unmarshal(&rawPC)
+	if err != nil {
+		return err
+	}
+
+	if *pc == nil {
+		*pc = make(map[string]PermissionLevel)
+	}
+	for key, value := range rawPC {
+		switch strings.ToLower(value) {
+		case "user":
+			(*pc)[key] = PermissionLevelUser
+		case "admin":
+			(*pc)[key] = PermissionLevelAdmin
+		default:
+			val, err := strconv.Atoi(value)
+			if err != nil {
+				(*pc)[key] = PermissionLevelDefault
+			} else {
+				(*pc)[key] = PermissionLevel(val)
+			}
+		}
+	}
+	return nil
+}
+
+func (pc *PermissionConfig) MarshalYAML() (interface{}, error) {
+	if *pc == nil {
+		return nil, nil
+	}
+	rawPC := make(map[string]string)
+	for key, value := range *pc {
+		switch value {
+		case PermissionLevelUser:
+			rawPC[key] = "user"
+		case PermissionLevelAdmin:
+			rawPC[key] = "admin"
+		default:
+			rawPC[key] = strconv.Itoa(int(value))
+		}
+	}
+	return rawPC, nil
+}
+
+func (pc PermissionConfig) IsWhitelisted(userID string) bool {
+	return pc.GetPermissionLevel(userID) >= 10
+}
+
+func (pc PermissionConfig) IsAdmin(userID string) bool {
+	return pc.GetPermissionLevel(userID) >= 100
+}
+
+func (pc PermissionConfig) GetPermissionLevel(userID string) PermissionLevel {
+	permissions, ok := pc[userID]
+	if ok {
+		return permissions
+	}
+
+	_, homeserver := appservice.ParseUserID(userID)
+	permissions, ok = pc[homeserver]
+	if len(homeserver) > 0 && ok {
+		return permissions
+	}
+
+	permissions, ok = pc["*"]
+	if ok {
+		return permissions
+	}
+
+	return PermissionLevelDefault
+}

+ 2 - 0
config/config.go

@@ -38,6 +38,8 @@ type Config struct {
 			URI  string `yaml:"uri"`
 		} `yaml:"database"`
 
+		StateStore string `yaml:"state_store_path"`
+
 		ID  string `yaml:"id"`
 		Bot struct {
 			Username    string `yaml:"username"`

+ 1 - 1
config/registration.go

@@ -54,7 +54,7 @@ func (config *Config) copyToRegistration(registration *appservice.Registration)
 	registration.RateLimited = false
 	registration.SenderLocalpart = config.AppService.Bot.Username
 
-	userIDRegex, err := regexp.Compile(fmt.Sprintf("@%s:%s",
+	userIDRegex, err := regexp.Compile(fmt.Sprintf("^@%s:%s$",
 		config.Bridge.FormatUsername("[0-9]+", "[0-9]+"),
 		config.Homeserver.Domain))
 	if err != nil {

+ 7 - 6
database/portal.go

@@ -18,6 +18,7 @@ package database
 
 import (
 	log "maunium.net/go/maulogger"
+	"maunium.net/go/mautrix-whatsapp/types"
 )
 
 type PortalQuery struct {
@@ -44,7 +45,7 @@ func (pq *PortalQuery) New() *Portal {
 	}
 }
 
-func (pq *PortalQuery) GetAll(owner string) (portals []*Portal) {
+func (pq *PortalQuery) GetAll(owner types.MatrixUserID) (portals []*Portal) {
 	rows, err := pq.db.Query("SELECT * FROM portal WHERE owner=?", owner)
 	if err != nil || rows == nil {
 		return nil
@@ -56,11 +57,11 @@ func (pq *PortalQuery) GetAll(owner string) (portals []*Portal) {
 	return
 }
 
-func (pq *PortalQuery) GetByJID(owner, jid string) *Portal {
+func (pq *PortalQuery) GetByJID(owner types.MatrixUserID, jid types.WhatsAppID) *Portal {
 	return pq.get("SELECT * FROM portal WHERE jid=? AND owner=?", jid, owner)
 }
 
-func (pq *PortalQuery) GetByMXID(mxid string) *Portal {
+func (pq *PortalQuery) GetByMXID(mxid types.MatrixRoomID) *Portal {
 	return pq.get("SELECT * FROM portal WHERE mxid=?", mxid)
 }
 
@@ -76,9 +77,9 @@ type Portal struct {
 	db  *Database
 	log log.Logger
 
-	JID   string
-	MXID  string
-	Owner string
+	JID   types.WhatsAppID
+	MXID  types.MatrixRoomID
+	Owner types.MatrixUserID
 }
 
 func (portal *Portal) Scan(row Scannable) *Portal {

+ 6 - 5
database/puppet.go

@@ -18,6 +18,7 @@ package database
 
 import (
 	log "maunium.net/go/maulogger"
+	"maunium.net/go/mautrix-whatsapp/types"
 )
 
 type PuppetQuery struct {
@@ -45,8 +46,8 @@ func (pq *PuppetQuery) New() *Puppet {
 	}
 }
 
-func (pq *PuppetQuery) GetAll() (puppets []*Puppet) {
-	rows, err := pq.db.Query("SELECT * FROM puppet")
+func (pq *PuppetQuery) GetAll(receiver types.MatrixUserID) (puppets []*Puppet) {
+	rows, err := pq.db.Query("SELECT * FROM puppet WHERE receiver=%s")
 	if err != nil || rows == nil {
 		return nil
 	}
@@ -57,7 +58,7 @@ func (pq *PuppetQuery) GetAll() (puppets []*Puppet) {
 	return
 }
 
-func (pq *PuppetQuery) Get(jid, receiver string) *Puppet {
+func (pq *PuppetQuery) Get(jid types.WhatsAppID, receiver types.MatrixUserID) *Puppet {
 	row := pq.db.QueryRow("SELECT * FROM user WHERE jid=? AND receiver=?", jid, receiver)
 	if row == nil {
 		return nil
@@ -69,8 +70,8 @@ type Puppet struct {
 	db  *Database
 	log log.Logger
 
-	JID      string
-	Receiver string
+	JID      types.WhatsAppID
+	Receiver types.MatrixUserID
 
 	Displayname string
 	Avatar      string

+ 4 - 3
database/user.go

@@ -19,6 +19,7 @@ package database
 import (
 	log "maunium.net/go/maulogger"
 	"github.com/Rhymen/go-whatsapp"
+	"maunium.net/go/mautrix-whatsapp/types"
 )
 
 type UserQuery struct {
@@ -61,7 +62,7 @@ func (uq *UserQuery) GetAll() (users []*User) {
 	return
 }
 
-func (uq *UserQuery) Get(userID string) *User {
+func (uq *UserQuery) Get(userID types.MatrixUserID) *User {
 	row := uq.db.QueryRow("SELECT * FROM user WHERE mxid=?", userID)
 	if row == nil {
 		return nil
@@ -73,8 +74,8 @@ type User struct {
 	db  *Database
 	log log.Logger
 
-	UserID         string
-	ManagementRoom string
+	UserID         types.MatrixUserID
+	ManagementRoom types.MatrixRoomID
 	Session        *whatsapp.Session
 }
 

+ 19 - 3
example-config.yaml

@@ -22,12 +22,15 @@ appservice:
     # The database URI. Usually file name. https://github.com/mattn/go-sqlite3#connection-string
     uri: mautrix-whatsapp.db
 
+  # Path to the Matrix room state store.
+  state_store_path: ./mx-state.json
+
   # The unique ID of this appservice.
   id: whatsapp
   # Appservice bot details.
   bot:
     # Username of the appservice bot.
-    username: whatsappbot
+    username: whatsapp
     # Display name and avatar for bot. Set to "remove" to remove display name/avatar, leave empty
     # to leave display name/avatar as-is.
     displayname: WhatsApp bridge bot
@@ -46,8 +49,21 @@ bridge:
   # Displayname template for WhatsApp users.
   # {{.displayname}} is replaced with the display name of the WhatsApp user.
   displayname_template: "{{.Displayname}}"
-  # Path to the Matrix room state store.
-  state_store_path: ./mx-state.json
+
+  # The prefix for commands. Only required in non-management rooms.
+  command_prefix: "!wa"
+
+  # Permissions for using the bridge.
+  # Permitted values:
+  #     user - Access to use the bridge to chat with a WhatsApp account.
+  #    admin - User level and some additional administration tools
+  # Permitted keys:
+  #        * - All Matrix users
+  #   domain - All users on that homeserver
+  #     mxid - Specific user
+  permissions:
+    "example.com": full
+    "@admin:example.com": admin
 
 # Logging config.
 logging:

+ 23 - 104
main.go

@@ -17,13 +17,8 @@
 package main
 
 import (
-	"github.com/Rhymen/go-whatsapp"
-	"time"
 	"fmt"
 	"os"
-	"bufio"
-	"encoding/gob"
-	"github.com/mdp/qrterminal"
 	"maunium.net/go/mautrix-whatsapp/config"
 	flag "maunium.net/go/mauflag"
 	"os/signal"
@@ -31,6 +26,8 @@ import (
 	"maunium.net/go/mautrix-appservice"
 	log "maunium.net/go/maulogger"
 	"maunium.net/go/mautrix-whatsapp/database"
+	"maunium.net/go/gomatrix"
+	"maunium.net/go/mautrix-whatsapp/types"
 )
 
 var configPath = flag.MakeFull("c", "config", "The path to your config file.", "config.yaml").String()
@@ -61,20 +58,21 @@ func (bridge *Bridge) GenerateRegistration() {
 }
 
 type Bridge struct {
-	AppService *appservice.AppService
-	Config     *config.Config
-	DB         *database.Database
-	Log        log.Logger
+	AppService     *appservice.AppService
+	EventProcessor *appservice.EventProcessor
+	Config         *config.Config
+	DB             *database.Database
+	Log            log.Logger
 
 	StateStore *AutosavingStateStore
 
-	MatrixListener *MatrixListener
-
-	users map[string]*User
+	users map[types.MatrixUserID]*User
 }
 
 func NewBridge() *Bridge {
-	bridge := &Bridge{}
+	bridge := &Bridge{
+		users: make(map[types.MatrixUserID]*User),
+	}
 	var err error
 	bridge.Config, err = config.Load(*configPath)
 	if err != nil {
@@ -97,7 +95,8 @@ func (bridge *Bridge) Init() {
 	log.DefaultLogger = bridge.Log.(*log.BasicLogger)
 	bridge.AppService.Log = log.Sub("Matrix")
 
-	bridge.StateStore = NewAutosavingStateStore(bridge.Config.Bridge.StateStore)
+	bridge.Log.Debugln("Initializing state store")
+	bridge.StateStore = NewAutosavingStateStore(bridge.Config.AppService.StateStore)
 	err = bridge.StateStore.Load()
 	if err != nil {
 		bridge.Log.Fatalln("Failed to load state store:", err)
@@ -105,19 +104,26 @@ func (bridge *Bridge) Init() {
 	}
 	bridge.AppService.StateStore = bridge.StateStore
 
+	bridge.Log.Debugln("Initializing database")
 	bridge.DB, err = database.New(bridge.Config.AppService.Database.URI)
 	if err != nil {
 		bridge.Log.Fatalln("Failed to initialize database:", err)
 		os.Exit(13)
 	}
 
-	bridge.MatrixListener = NewMatrixListener(bridge)
+	bridge.Log.Debugln("Initializing event processor")
+	bridge.EventProcessor = appservice.NewEventProcessor(bridge.AppService)
+	bridge.EventProcessor.On(gomatrix.EventMessage, bridge.HandleMessage)
+	bridge.EventProcessor.On(gomatrix.StateMember, bridge.HandleMembership)
 }
 
 func (bridge *Bridge) Start() {
 	bridge.DB.CreateTables()
+	bridge.Log.Debugln("Starting application service HTTP server")
 	go bridge.AppService.Start()
-	go bridge.MatrixListener.Start()
+	bridge.Log.Debugln("Starting event processor")
+	go bridge.EventProcessor.Start()
+	bridge.Log.Debugln("Updating bot profile")
 	go bridge.UpdateBotProfile()
 }
 
@@ -146,7 +152,7 @@ func (bridge *Bridge) UpdateBotProfile() {
 
 func (bridge *Bridge) Stop() {
 	bridge.AppService.Stop()
-	bridge.MatrixListener.Stop()
+	bridge.EventProcessor.Stop()
 	err := bridge.StateStore.Save()
 	if err != nil {
 		bridge.Log.Warnln("Failed to save state store:", err)
@@ -190,90 +196,3 @@ func main() {
 
 	NewBridge().Main()
 }
-
-func temp() {
-	wac, err := whatsapp.NewConn(20 * time.Second)
-	if err != nil {
-		panic(err)
-	}
-
-	wac.AddHandler(myHandler{})
-
-	sess, err := LoadSession("whatsapp.session")
-	if err != nil {
-		fmt.Println(err)
-		sess, err = Login(wac)
-	} else {
-		sess, err = wac.RestoreSession(sess)
-	}
-	if err != nil {
-		panic(err)
-	}
-	SaveSession(sess, "whatsapp.session")
-
-	reader := bufio.NewReader(os.Stdin)
-	for {
-		fmt.Print("receiver> ")
-		receiver, _ := reader.ReadString('\n')
-		fmt.Print("message> ")
-		message, _ := reader.ReadString('\n')
-		wac.Send(whatsapp.TextMessage{
-			Info: whatsapp.MessageInfo{
-				RemoteJid: fmt.Sprintf("%s@s.whatsapp.net", receiver),
-			},
-			Text: message,
-		})
-		fmt.Println(receiver, message)
-	}
-}
-
-func Login(wac *whatsapp.Conn) (whatsapp.Session, error) {
-	qrChan := make(chan string)
-	go func() {
-		qrterminal.Generate(<-qrChan, qrterminal.L, os.Stdout)
-	}()
-	return wac.Login(qrChan)
-}
-
-func SaveSession(session whatsapp.Session, fileName string) {
-	file, err := os.OpenFile(fileName, os.O_WRONLY|os.O_CREATE, 0600)
-	if err != nil {
-		panic(err)
-	}
-
-	enc := gob.NewEncoder(file)
-	enc.Encode(session)
-}
-
-func LoadSession(fileName string) (sess whatsapp.Session, err error) {
-	file, err := os.OpenFile(fileName, os.O_RDONLY, 0600)
-	if err != nil {
-		return sess, err
-	}
-
-	dec := gob.NewDecoder(file)
-	dec.Decode(sess)
-	return
-}
-
-type myHandler struct{}
-
-func (myHandler) HandleError(err error) {
-	fmt.Fprintf(os.Stderr, "%v", err)
-}
-
-func (myHandler) HandleTextMessage(message whatsapp.TextMessage) {
-	fmt.Println(message)
-}
-
-func (myHandler) HandleImageMessage(message whatsapp.ImageMessage) {
-	fmt.Println(message)
-}
-
-func (myHandler) HandleVideoMessage(message whatsapp.VideoMessage) {
-	fmt.Println(message)
-}
-
-func (myHandler) HandleJsonMessage(message string) {
-	fmt.Println(message)
-}

+ 22 - 58
matrix.go

@@ -17,96 +17,60 @@
 package main
 
 import (
-	log "maunium.net/go/maulogger"
-	"maunium.net/go/mautrix-appservice"
 	"maunium.net/go/gomatrix"
 )
 
-type MatrixListener struct {
-	bridge *Bridge
-	as     *appservice.AppService
-	log    log.Logger
-	stop   chan struct{}
-}
-
-func NewMatrixListener(bridge *Bridge) *MatrixListener {
-	return &MatrixListener{
-		bridge: bridge,
-		as:     bridge.AppService,
-		stop:   make(chan struct{}, 1),
-		log:    bridge.Log.Sub("Matrix Listener"),
-	}
-}
-
-func (ml *MatrixListener) Start() {
-	for {
-		select {
-		case evt := <-ml.bridge.AppService.Events:
-			ml.log.Debugln("Received Matrix event:", evt)
-			switch evt.Type {
-			case gomatrix.StateMember:
-				ml.HandleMembership(evt)
-			case gomatrix.EventMessage:
-				ml.HandleMessage(evt)
-			}
-		case <-ml.stop:
-			return
-		}
-	}
-}
-
-func (ml *MatrixListener) HandleBotInvite(evt *gomatrix.Event) {
-	intent := ml.as.BotIntent()
+func (bridge *Bridge) HandleBotInvite(evt *gomatrix.Event) {
+	intent := bridge.AppService.BotIntent()
 
 	resp, err := intent.JoinRoom(evt.RoomID, "", nil)
 	if err != nil {
-		ml.log.Debugln("Failed to join room", evt.RoomID, "with invite from", evt.Sender)
+		bridge.Log.Debugln("Failed to join room", evt.RoomID, "with invite from", evt.Sender)
 		return
 	}
 
 	members, err := intent.JoinedMembers(resp.RoomID)
 	if err != nil {
-		ml.log.Debugln("Failed to get members in room", resp.RoomID, "after accepting invite from", evt.Sender)
+		bridge.Log.Debugln("Failed to get members in room", resp.RoomID, "after accepting invite from", evt.Sender)
 		intent.LeaveRoom(resp.RoomID)
 		return
 	}
 
 	if len(members.Joined) < 2 {
-		ml.log.Debugln("Leaving empty room", resp.RoomID, "after accepting invite from", evt.Sender)
+		bridge.Log.Debugln("Leaving empty room", resp.RoomID, "after accepting invite from", evt.Sender)
 		intent.LeaveRoom(resp.RoomID)
 		return
 	}
+
+	hasPuppets := false
 	for mxid, _ := range members.Joined {
 		if mxid == intent.UserID || mxid == evt.Sender {
 			continue
-		} else if true { // TODO check if mxid is WhatsApp puppet
-
+		} else if _, _, ok := bridge.ParsePuppetMXID(mxid); ok {
+			hasPuppets = true
 			continue
 		}
-		ml.log.Debugln("Leaving multi-user room", resp.RoomID, "after accepting invite from", evt.Sender)
+		bridge.Log.Debugln("Leaving multi-user room", resp.RoomID, "after accepting invite from", evt.Sender)
 		intent.SendNotice(resp.RoomID, "This bridge is user-specific, please don't invite me into rooms with other users.")
 		intent.LeaveRoom(resp.RoomID)
 		return
 	}
 
-	user := ml.bridge.GetUser(evt.Sender)
-	user.ManagementRoom = resp.RoomID
-	user.Update()
-	intent.SendNotice(user.ManagementRoom, "This room has been registered as your bridge management/status room.")
-	ml.log.Debugln(resp.RoomID, "registered as a management room with", evt.Sender)
-}
-
-func (ml *MatrixListener) HandleMembership(evt *gomatrix.Event) {
-	ml.log.Debugln(evt.Content, evt.Content.Membership, evt.GetStateKey())
-	if evt.Content.Membership == "invite" && evt.GetStateKey() == ml.as.BotMXID() {
-		ml.HandleBotInvite(evt)
+	if !hasPuppets {
+		user := bridge.GetUser(evt.Sender)
+		user.ManagementRoom = resp.RoomID
+		user.Update()
+		intent.SendNotice(user.ManagementRoom, "This room has been registered as your bridge management/status room.")
+		bridge.Log.Debugln(resp.RoomID, "registered as a management room with", evt.Sender)
 	}
 }
 
-func (ml *MatrixListener) HandleMessage(evt *gomatrix.Event) {
-
+func (bridge *Bridge) HandleMembership(evt *gomatrix.Event) {
+	bridge.Log.Debugln(evt.Content, evt.Content.Membership, evt.GetStateKey())
+	if evt.Content.Membership == "invite" && evt.GetStateKey() == bridge.AppService.BotMXID() {
+		bridge.HandleBotInvite(evt)
+	}
 }
 
-func (ml *MatrixListener) Stop() {
-	ml.stop <- struct{}{}
+func (bridge *Bridge) HandleMessage(evt *gomatrix.Event) {
 }

+ 3 - 2
portal.go

@@ -20,9 +20,10 @@ import (
 	"maunium.net/go/mautrix-whatsapp/database"
 	log "maunium.net/go/maulogger"
 	"fmt"
+	"maunium.net/go/mautrix-whatsapp/types"
 )
 
-func (user *User) GetPortalByMXID(mxid string) *Portal {
+func (user *User) GetPortalByMXID(mxid types.MatrixRoomID) *Portal {
 	portal, ok := user.portalsByMXID[mxid]
 	if !ok {
 		dbPortal := user.bridge.DB.Portal.GetByMXID(mxid)
@@ -38,7 +39,7 @@ func (user *User) GetPortalByMXID(mxid string) *Portal {
 	return portal
 }
 
-func (user *User) GetPortalByJID(jid string) *Portal {
+func (user *User) GetPortalByJID(jid types.WhatsAppID) *Portal {
 	portal, ok := user.portalsByJID[jid]
 	if !ok {
 		dbPortal := user.bridge.DB.Portal.GetByJID(user.UserID, jid)

+ 113 - 0
puppet.go

@@ -0,0 +1,113 @@
+// mautrix-whatsapp - A Matrix-WhatsApp puppeting bridge.
+// Copyright (C) 2018 Tulir Asokan
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+package main
+
+import (
+	"maunium.net/go/mautrix-whatsapp/database"
+	log "maunium.net/go/maulogger"
+	"fmt"
+	"regexp"
+	"maunium.net/go/mautrix-whatsapp/types"
+	"strings"
+)
+
+func (bridge *Bridge) ParsePuppetMXID(mxid types.MatrixUserID) (types.MatrixUserID, types.WhatsAppID, bool) {
+	userIDRegex, err := regexp.Compile(fmt.Sprintf("^@%s:%s$",
+		bridge.Config.Bridge.FormatUsername("([0-9]+)", "([0-9]+)"),
+		bridge.Config.Homeserver.Domain))
+	if err != nil {
+		bridge.Log.Warnln("Failed to compile puppet user ID regex:", err)
+		return "", "", false
+	}
+	match := userIDRegex.FindStringSubmatch(string(mxid))
+	if match == nil || len(match) != 3 {
+		return "", "", false
+	}
+
+	receiver := match[1]
+	receiver = strings.Replace(receiver, "=40", "@", 1)
+	colonIndex := strings.LastIndex(receiver, "=3")
+	receiver = receiver[:colonIndex] + ":" + receiver[colonIndex+len("=3"):]
+	return types.MatrixUserID(receiver), types.WhatsAppID(match[2]), true
+}
+
+func (bridge *Bridge) GetPuppetByMXID(mxid types.MatrixUserID) *Puppet {
+	receiver, jid, ok := bridge.ParsePuppetMXID(mxid)
+	if !ok {
+		return nil
+	}
+
+	user := bridge.GetUser(receiver)
+	if user == nil {
+		return nil
+	}
+
+	return user.GetPuppetByJID(jid)
+}
+
+func (user *User) GetPuppetByMXID(mxid types.MatrixUserID) *Puppet {
+	receiver, jid, ok := user.bridge.ParsePuppetMXID(mxid)
+	if !ok || receiver != user.UserID {
+		return nil
+	}
+
+	return user.GetPuppetByJID(jid)
+}
+
+func (user *User) GetPuppetByJID(jid types.WhatsAppID) *Puppet {
+	puppet, ok := user.puppets[jid]
+	if !ok {
+		dbPuppet := user.bridge.DB.Puppet.Get(jid, user.UserID)
+		if dbPuppet == nil {
+			return nil
+		}
+		puppet = user.NewPuppet(dbPuppet)
+		user.puppets[puppet.JID] = puppet
+	}
+	return puppet
+}
+
+func (user *User) GetAllPuppets() []*Puppet {
+	dbPuppets := user.bridge.DB.Puppet.GetAll(user.UserID)
+	output := make([]*Puppet, len(dbPuppets))
+	for index, dbPuppet := range dbPuppets {
+		puppet, ok := user.puppets[dbPuppet.JID]
+		if !ok {
+			puppet = user.NewPuppet(dbPuppet)
+			user.puppets[dbPuppet.JID] = puppet
+		}
+		output[index] = puppet
+	}
+	return output
+}
+
+func (user *User) NewPuppet(dbPuppet *database.Puppet) *Puppet {
+	return &Puppet{
+		Puppet: dbPuppet,
+		user:   user,
+		bridge: user.bridge,
+		log:    user.log.Sub(fmt.Sprintf("Puppet/%s", dbPuppet.JID)),
+	}
+}
+
+type Puppet struct {
+	*database.Puppet
+
+	user   *User
+	bridge *Bridge
+	log    log.Logger
+}

+ 26 - 0
types/types.go

@@ -0,0 +1,26 @@
+// mautrix-whatsapp - A Matrix-WhatsApp puppeting bridge.
+// Copyright (C) 2018 Tulir Asokan
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+package types
+
+// WhatsAppID is a WhatsApp JID.
+type WhatsAppID = string
+
+// MatrixUserID is the ID of a Matrix user.
+type MatrixUserID = string
+
+// MatrixRoomID is the internal room ID of a Matrix room.
+type MatrixRoomID = string

+ 5 - 4
user.go

@@ -24,6 +24,7 @@ import (
 	"os"
 	"github.com/skip2/go-qrcode"
 	log "maunium.net/go/maulogger"
+	"maunium.net/go/mautrix-whatsapp/types"
 )
 
 type User struct {
@@ -33,12 +34,12 @@ type User struct {
 	bridge *Bridge
 	log    log.Logger
 
-	portalsByMXID map[string]*Portal
-	portalsByJID  map[string]*Portal
-	puppets       map[string]*Portal
+	portalsByMXID map[types.MatrixRoomID]*Portal
+	portalsByJID  map[types.WhatsAppID]*Portal
+	puppets       map[types.WhatsAppID]*Puppet
 }
 
-func (bridge *Bridge) GetUser(userID string) *User {
+func (bridge *Bridge) GetUser(userID types.MatrixUserID) *User {
 	user, ok := bridge.users[userID]
 	if !ok {
 		dbUser := bridge.DB.User.Get(userID)