Fix data storage and other things

This commit is contained in:
Tulir Asokan 2018-08-16 19:20:07 +03:00
parent 9c48eeb534
commit 141eba644b
13 changed files with 158 additions and 36 deletions

5
.gitignore vendored
View File

@ -1,5 +1,8 @@
.idea
*.session
*.yaml
!example-config.yaml
*.session
*.json
*.db

View File

@ -24,6 +24,7 @@ import (
type BridgeConfig struct {
UsernameTemplate string `yaml:"username_template"`
DisplaynameTemplate string `yaml:"displayname_template"`
StateStore string `yaml:"state_store_path"`
usernameTemplate *template.Template `yaml:"-"`
displaynameTemplate *template.Template `yaml:"-"`
}

View File

@ -78,6 +78,8 @@ func (config *Config) MakeAppService() (*appservice.AppService, error) {
as.LogConfig = config.Logging
as.HomeserverDomain = config.Homeserver.Domain
as.HomeserverURL = config.Homeserver.Address
as.Host.Hostname = config.AppService.Hostname
as.Host.Port = config.AppService.Port
var err error
as.Registration, err = config.GetRegistration()
return as, err

View File

@ -24,7 +24,7 @@ import (
type Database struct {
*sql.DB
log *log.Sublogger
log log.Logger
User *UserQuery
Portal *PortalQuery
@ -39,23 +39,29 @@ func New(file string) (*Database, error) {
db := &Database{
DB: conn,
log: log.CreateSublogger("Database", log.LevelDebug),
log: log.Sub("Database"),
}
db.User = &UserQuery{
db: db,
log: log.CreateSublogger("Database/User", log.LevelDebug),
log: db.log.Sub("User"),
}
db.Portal = &PortalQuery{
db: db,
log: log.CreateSublogger("Database/Portal", log.LevelDebug),
log: db.log.Sub("Portal"),
}
db.Puppet = &PuppetQuery{
db: db,
log: log.CreateSublogger("Database/Puppet", log.LevelDebug),
log: db.log.Sub("Puppet"),
}
return db, nil
}
func (db *Database) CreateTables() {
db.User.CreateTable()
db.Portal.CreateTable()
db.Puppet.CreateTable()
}
type Scannable interface {
Scan(...interface{}) error
}

View File

@ -22,7 +22,7 @@ import (
type PortalQuery struct {
db *Database
log *log.Sublogger
log log.Logger
}
func (pq *PortalQuery) CreateTable() error {
@ -74,7 +74,7 @@ func (pq *PortalQuery) get(query string, args ...interface{}) *Portal {
type Portal struct {
db *Database
log *log.Sublogger
log log.Logger
JID string
MXID string

View File

@ -22,7 +22,7 @@ import (
type PuppetQuery struct {
db *Database
log *log.Sublogger
log log.Logger
}
func (pq *PuppetQuery) CreateTable() error {
@ -67,7 +67,7 @@ func (pq *PuppetQuery) Get(jid, receiver string) *Puppet {
type Puppet struct {
db *Database
log *log.Sublogger
log log.Logger
JID string
Receiver string

View File

@ -23,7 +23,7 @@ import (
type UserQuery struct {
db *Database
log *log.Sublogger
log log.Logger
}
func (uq *UserQuery) CreateTable() error {
@ -71,7 +71,7 @@ func (uq *UserQuery) Get(userID string) *User {
type User struct {
db *Database
log *log.Sublogger
log log.Logger
UserID string
ManagementRoom string

View File

@ -46,6 +46,8 @@ 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
# Logging config.
logging:
@ -61,4 +63,4 @@ logging:
timestamp_format: Jan _2, 2006 15:04:05
# Minimum severity for log messages.
# Options: debug, info, warn, error, fatal
print_level: info
print_level: debug

54
main.go
View File

@ -64,7 +64,9 @@ type Bridge struct {
AppService *appservice.AppService
Config *config.Config
DB *database.Database
Log *log.Logger
Log log.Logger
StateStore *AutosavingStateStore
MatrixListener *MatrixListener
@ -84,33 +86,71 @@ func NewBridge() *Bridge {
func (bridge *Bridge) Init() {
var err error
bridge.AppService, err = bridge.Config.MakeAppService()
if err != nil {
fmt.Fprintln(os.Stderr, "Failed to initialize AppService:", err)
os.Exit(11)
}
bridge.AppService.Init()
bridge.Log = bridge.AppService.Log.Parent
log.DefaultLogger = bridge.Log
bridge.AppService.Log = log.CreateSublogger("Matrix", log.LevelDebug)
bridge.Log = bridge.AppService.Log
log.DefaultLogger = bridge.Log.(*log.BasicLogger)
bridge.AppService.Log = log.Sub("Matrix")
bridge.StateStore = NewAutosavingStateStore(bridge.Config.Bridge.StateStore)
err = bridge.StateStore.Load()
if err != nil {
bridge.Log.Fatalln("Failed to load state store:", err)
os.Exit(12)
}
bridge.AppService.StateStore = bridge.StateStore
bridge.DB, err = database.New(bridge.Config.AppService.Database.URI)
if err != nil {
bridge.Log.Fatalln("Failed to initialize database:", err)
os.Exit(12)
os.Exit(13)
}
bridge.MatrixListener = NewMatrixListener(bridge)
}
func (bridge *Bridge) Start() {
bridge.AppService.Start()
bridge.MatrixListener.Start()
bridge.DB.CreateTables()
go bridge.AppService.Start()
go bridge.MatrixListener.Start()
go bridge.UpdateBotProfile()
}
func (bridge *Bridge) UpdateBotProfile() {
botConfig := bridge.Config.AppService.Bot
var err error
if botConfig.Avatar == "remove" {
err = bridge.AppService.BotIntent().SetAvatarURL("")
} else if len(botConfig.Avatar) > 0 {
err = bridge.AppService.BotIntent().SetAvatarURL(botConfig.Avatar)
}
if err != nil {
bridge.Log.Warnln("Failed to update bot avatar:", err)
}
if botConfig.Displayname == "remove" {
err = bridge.AppService.BotIntent().SetDisplayName("")
} else if len(botConfig.Avatar) > 0 {
err = bridge.AppService.BotIntent().SetDisplayName(botConfig.Displayname)
}
if err != nil {
bridge.Log.Warnln("Failed to update bot displayname:", err)
}
}
func (bridge *Bridge) Stop() {
bridge.AppService.Stop()
bridge.MatrixListener.Stop()
err := bridge.StateStore.Save()
if err != nil {
bridge.Log.Warnln("Failed to save state store:", err)
}
}
func (bridge *Bridge) Main() {

View File

@ -25,7 +25,7 @@ import (
type MatrixListener struct {
bridge *Bridge
as *appservice.AppService
log *log.Sublogger
log log.Logger
stop chan struct{}
}
@ -34,7 +34,7 @@ func NewMatrixListener(bridge *Bridge) *MatrixListener {
bridge: bridge,
as: bridge.AppService,
stop: make(chan struct{}, 1),
log: bridge.Log.CreateSublogger("Matrix", log.LevelDebug),
log: bridge.Log.Sub("Matrix Listener"),
}
}
@ -42,7 +42,7 @@ func (ml *MatrixListener) Start() {
for {
select {
case evt := <-ml.bridge.AppService.Events:
log.Debugln("Received Matrix event:", evt)
ml.log.Debugln("Received Matrix event:", evt)
switch evt.Type {
case gomatrix.StateMember:
ml.HandleMembership(evt)
@ -56,47 +56,48 @@ func (ml *MatrixListener) Start() {
}
func (ml *MatrixListener) HandleBotInvite(evt *gomatrix.Event) {
cli := ml.as.BotClient()
intent := ml.as.BotIntent()
resp, err := cli.JoinRoom(evt.RoomID, "", nil)
resp, err := intent.JoinRoom(evt.RoomID, "", nil)
if err != nil {
ml.log.Debugln("Failed to join room", evt.RoomID, "with invite from", evt.Sender)
return
}
members, err := cli.JoinedMembers(resp.RoomID)
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)
cli.LeaveRoom(resp.RoomID)
intent.LeaveRoom(resp.RoomID)
return
}
if len(members.Joined) < 2 {
ml.log.Debugln("Leaving empty room", resp.RoomID, "after accepting invite from", evt.Sender)
cli.LeaveRoom(resp.RoomID)
intent.LeaveRoom(resp.RoomID)
return
}
for mxid, _ := range members.Joined {
if mxid == cli.UserID || mxid == evt.Sender {
if mxid == intent.UserID || mxid == evt.Sender {
continue
} else if true { // TODO check if mxid is WhatsApp puppet
continue
}
ml.log.Debugln("Leaving multi-user room", resp.RoomID, "after accepting invite from", evt.Sender)
cli.SendNotice(resp.RoomID, "This bridge is user-specific, please don't invite me into rooms with other users.")
cli.LeaveRoom(resp.RoomID)
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()
cli.SendNotice(user.ManagementRoom, "This room has been registered as your bridge management/status room.")
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)
}

View File

@ -76,7 +76,7 @@ func (user *User) NewPortal(dbPortal *database.Portal) *Portal {
Portal: dbPortal,
user: user,
bridge: user.bridge,
log: user.bridge.Log.CreateSublogger(fmt.Sprintf("Portal/%s/%s", user.UserID, dbPortal.JID), log.LevelDebug),
log: user.log.Sub(fmt.Sprintf("Portal/%s", dbPortal.JID)),
}
}
@ -85,5 +85,5 @@ type Portal struct {
user *User
bridge *Bridge
log *log.Sublogger
log log.Logger
}

67
statestore.go Normal file
View File

@ -0,0 +1,67 @@
// 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-appservice"
"encoding/json"
"io/ioutil"
"os"
)
type AutosavingStateStore struct {
*appservice.BasicStateStore
Path string
}
func NewAutosavingStateStore(path string) *AutosavingStateStore {
return &AutosavingStateStore{
BasicStateStore: appservice.NewBasicStateStore(),
Path: path,
}
}
func (store *AutosavingStateStore) Save() error {
data, err := json.Marshal(store.BasicStateStore)
if err != nil {
return err
}
return ioutil.WriteFile(store.Path, data, 0600)
}
func (store *AutosavingStateStore) Load() error {
data, err := ioutil.ReadFile(store.Path)
if err != nil {
if os.IsNotExist(err) {
return nil
}
return err
}
return json.Unmarshal(data, store.BasicStateStore)
}
func (store *AutosavingStateStore) MarkRegistered(userID string) {
store.BasicStateStore.MarkRegistered(userID)
store.Save()
}
func (store *AutosavingStateStore) SetMembership(roomID, userID, membership string) {
store.BasicStateStore.SetMembership(roomID, userID, membership)
store.Save()
}

View File

@ -31,7 +31,7 @@ type User struct {
Conn *whatsapp.Conn
bridge *Bridge
log *log.Sublogger
log log.Logger
portalsByMXID map[string]*Portal
portalsByJID map[string]*Portal
@ -77,7 +77,7 @@ func (bridge *Bridge) NewUser(dbUser *database.User) *User {
return &User{
User: dbUser,
bridge: bridge,
log: bridge.Log.CreateSublogger(fmt.Sprintf("User/%s", dbUser.UserID), log.LevelDebug),
log: bridge.Log.Sub("User").Sub(dbUser.UserID),
}
}