misc/infrastructure/powerbot/bot.go

   1 package main
   2
   3 import (
   4         "encoding/json"
   5         "fmt"
   6         "io/ioutil"
   7         "log"
   8         "maunium.net/go/mautrix"
   9         "maunium.net/go/mautrix/event"
  10         "maunium.net/go/mautrix/id"
  11         "strings"
  12         "sync"
  13         "time"
  14 )
  15
  16 const (
  17         syncInterval       = time.Minute
  18         syncForceFrequency = int(7 * 24 * time.Hour / syncInterval)
  19 )
  20
  21 type Room struct {
  22         ID   id.RoomID
  23         Name string
  24 }
  25
  26 type Config struct {
  27         Homeserver  string      `json:"homeserver"`
  28         UserID      id.UserID   `json:"user_id"`
  29         Password    string      `json:"password,omitempty"`
  30         DeviceID    id.DeviceID `json:"device_id,omitempty"`
  31         AccessToken string      `json:"access_token,omitempty"`
  32         Rooms       [][]Room    `json:"rooms"`
  33 }
  34
  35 func (c *Config) Load() error {
  36         log.Printf("Loading config.")
  37         data, err := ioutil.ReadFile("config.json")
  38         if err != nil {
  39                 return err
  40         }
  41         return json.Unmarshal(data, c)
  42 }
  43
  44 func (c *Config) Save() error {
  45         log.Printf("Saving config.")
  46         data, err := json.MarshalIndent(c, "", "\t")
  47         if err != nil {
  48                 return err
  49         }
  50         return ioutil.WriteFile("config.json", data, 0700)
  51 }
  52
  53 func Login(config *Config) (*mautrix.Client, error) {
  54         configMu.Lock()
  55         defer configMu.Unlock()
  56
  57         // Note: we have to lower case the user ID for Matrix protocol communication.
  58         uid := id.UserID(strings.ToLower(string(config.UserID)))
  59         client, err := mautrix.NewClient(config.Homeserver, uid, config.AccessToken)
  60         if err != nil {
  61                 return nil, fmt.Errorf("failed to create client: %v", err)
  62         }
  63         if config.AccessToken == "" {
  64                 resp, err := client.Login(&mautrix.ReqLogin{
  65                         Type: mautrix.AuthTypePassword,
  66                         Identifier: mautrix.UserIdentifier{
  67                                 Type: mautrix.IdentifierTypeUser,
  68                                 User: string(client.UserID),
  69                         },
  70                         Password:                 config.Password,
  71                         InitialDeviceDisplayName: "matrixbot",
  72                         StoreCredentials:         true,
  73                 })
  74                 if err != nil {
  75                         return nil, fmt.Errorf("failed to authenticate: %v", err)
  76                 }
  77                 config.Password = ""
  78                 config.DeviceID = resp.DeviceID
  79                 config.AccessToken = resp.AccessToken
  80                 err = config.Save()
  81                 if err != nil {
  82                         return nil, fmt.Errorf("failed to save config: %v", err)
  83                 }
  84         } else {
  85                 client.DeviceID = config.DeviceID
  86         }
  87         return client, nil
  88 }
  89
  90 var (
  91         configMu sync.Mutex
  92
  93         roomUsersMu sync.RWMutex
  94         roomUsers   = map[id.RoomID]map[id.UserID]struct{}{}
  95
  96         fullySynced     bool
  97         roomPowerLevels = map[id.RoomID]*event.PowerLevelsEventContent{}
  98 )
  99
 100 func setUserStateAt(room id.RoomID, user id.UserID, now time.Time, maxPrevState, state State) {
 101         err := writeUserStateAt(room, user, now, maxPrevState, state)
 102         if err != nil {
 103                 log.Fatalf("failed to write user state: %v", err)
 104         }
 105 }
 106
 107 func handleMessage(now time.Time, room id.RoomID, sender id.UserID, raw *event.Event) {
 108         // log.Printf("[%v] Message from %v to %v", now, sender, room)
 109         roomUsersMu.Lock()
 110         roomUsers[room][sender] = struct{}{}
 111         roomUsersMu.Unlock()
 112         setUserStateAt(room, sender, now.Add(-activeTime), Active, Active)
 113         setUserStateAt(room, sender, now, Active, Idle)
 114 }
 115
 116 func handleJoin(now time.Time, room id.RoomID, member id.UserID, raw *event.Event) {
 117         log.Printf("[%v] Join from %v to %v", now, member, room)
 118         roomUsersMu.Lock()
 119         roomUsers[room][member] = struct{}{}
 120         roomUsersMu.Unlock()
 121         setUserStateAt(room, member, now, NotActive, Idle)
 122 }
 123
 124 func handleLeave(now time.Time, room id.RoomID, member id.UserID, raw *event.Event) {
 125         log.Printf("[%v] Leave from %v to %v", now, member, room)
 126         roomUsersMu.Lock()
 127         delete(roomUsers[room], member)
 128         roomUsersMu.Unlock()
 129         setUserStateAt(room, member, now, Active, NotActive)
 130 }
 131
 132 func handlePowerLevels(now time.Time, room id.RoomID, levels *event.PowerLevelsEventContent, raw *event.Event) {
 133         // log.Printf("[%v] Power levels for %v are %v", now, room, levels)
 134         levelsCopy := *levels // Looks like mautrix always passes the same pointer here.
 135         roomUsersMu.Lock()
 136         roomPowerLevels[room] = &levelsCopy
 137         roomUsersMu.Unlock()
 138 }
 139
 140 func eventTime(evt *event.Event) time.Time {
 141         return time.Unix(0, evt.Timestamp*1000000)
 142 }
 143
 144 type MoreMessagesSyncer struct {
 145         *mautrix.DefaultSyncer
 146 }
 147
 148 func newSyncer() *MoreMessagesSyncer {
 149         return &MoreMessagesSyncer{
 150                 DefaultSyncer: mautrix.NewDefaultSyncer(),
 151         }
 152 }
 153
 154 func (s *MoreMessagesSyncer) GetFilterJSON(userID id.UserID) *mautrix.Filter {
 155         f := s.DefaultSyncer.GetFilterJSON(userID)
 156         // Same filters as Element.
 157         f.Room.Timeline.Limit = 20
 158         // Only include our rooms.
 159         f.Room.Rooms = make([]id.RoomID, 0, len(roomUsers))
 160         for room := range roomUsers {
 161                 f.Room.Rooms = append(f.Room.Rooms, room)
 162         }
 163         return f
 164 }
 165
 166 func isRoom(room id.RoomID) bool {
 167         roomUsersMu.RLock()
 168         defer roomUsersMu.RUnlock()
 169         _, found := roomUsers[room]
 170         return found
 171 }
 172
 173 func Run() (err error) {
 174         err = InitDatabase()
 175         if err != nil {
 176                 return fmt.Errorf("failed to init database: %v", err)
 177         }
 178         defer func() {
 179                 err2 := CloseDatabase()
 180                 if err2 != nil && err == nil {
 181                         err = fmt.Errorf("failed to close database: %v", err)
 182                 }
 183         }()
 184         logPowerLevelBounds()
 185         config := &Config{}
 186         err = config.Load()
 187         if err != nil {
 188                 return fmt.Errorf("failed to load config: %v", err)
 189         }
 190         for _, group := range config.Rooms {
 191                 for _, room := range group {
 192                         roomUsers[room.ID] = map[id.UserID]struct{}{}
 193                 }
 194         }
 195         client, err := Login(config)
 196         if err != nil {
 197                 return fmt.Errorf("failed to login: %v", err)
 198         }
 199         syncer := newSyncer()
 200         syncer.OnEventType(event.StateTombstone, func(source mautrix.EventSource, evt *event.Event) {
 201                 if !isRoom(evt.RoomID) {
 202                         return
 203                 }
 204                 tomb := evt.Content.AsTombstone()
 205                 if tomb.ReplacementRoom == "" {
 206                         log.Printf("Replacement room in tombstone event is not set - not handling: %v", evt)
 207                         return
 208                 }
 209                 for _, group := range config.Rooms {
 210                         for _, room := range group {
 211                                 if room.ID == evt.RoomID {
 212                                         configMu.Lock()
 213                                         defer configMu.Unlock()
 214                                         room.ID = tomb.ReplacementRoom
 215                                         config.Save()
 216                                         log.Fatalf("room upgrade handled from %v to %v - need restart", evt.RoomID, tomb.ReplacementRoom)
 217                                 }
 218                         }
 219                 }
 220         })
 221         syncer.OnEventType(event.EventMessage, func(source mautrix.EventSource, evt *event.Event) {
 222                 if !isRoom(evt.RoomID) {
 223                         return
 224                 }
 225                 handleMessage(eventTime(evt), evt.RoomID, evt.Sender, evt)
 226         })
 227         syncer.OnEventType(event.StateMember, func(source mautrix.EventSource, evt *event.Event) {
 228                 if !isRoom(evt.RoomID) {
 229                         return
 230                 }
 231                 mem := evt.Content.AsMember()
 232                 key := evt.StateKey
 233                 if key == nil {
 234                         return
 235                 }
 236                 member := id.UserID(*key)
 237                 switch mem.Membership {
 238                 case event.MembershipJoin:
 239                         handleJoin(eventTime(evt), evt.RoomID, member, evt)
 240                 case event.MembershipLeave, event.MembershipBan:
 241                         handleLeave(eventTime(evt), evt.RoomID, member, evt)
 242                 default: // Ignore.
 243                 }
 244         })
 245         syncer.OnEventType(event.StatePowerLevels, func(source mautrix.EventSource, evt *event.Event) {
 246                 if !isRoom(evt.RoomID) {
 247                         return
 248                 }
 249                 handlePowerLevels(eventTime(evt), evt.RoomID, evt.Content.AsPowerLevels(), evt)
 250         })
 251         syncer.OnSync(func(resp *mautrix.RespSync, since string) bool {
 252                 // j, _ := json.MarshalIndent(resp, "", "  ")
 253                 // log.Print(string(j))
 254                 roomUsersMu.Lock()
 255                 if since != "" && !fullySynced {
 256                         log.Print("Fully synced.")
 257                         for room, users := range roomUsers {
 258                                 if _, found := users[config.UserID]; !found {
 259                                         log.Printf("Not actually joined %v yet...", room)
 260                                         _, err := client.JoinRoom(string(room), "", nil)
 261                                         if err != nil {
 262                                                 log.Printf("Failed to join %v: %v", room, err)
 263                                         }
 264                                 }
 265                         }
 266                         fullySynced = true
 267                 }
 268                 roomUsersMu.Unlock()
 269                 return true
 270         })
 271         client.Syncer = syncer
 272         ticker := time.NewTicker(syncInterval)
 273         defer ticker.Stop()
 274         go func() {
 275                 counter := 0
 276                 for range ticker.C {
 277                         roomUsersMu.RLock()
 278                         scoreData := map[id.RoomID]map[id.UserID]*Score{}
 279                         now := time.Now()
 280                         for room := range roomUsers {
 281                                 scores, err := queryUserScores(room, now)
 282                                 if err != nil {
 283                                         log.Fatalf("failed to query user scores: %v", err)
 284                                 }
 285                                 scoreData[room] = scores
 286                         }
 287                         for _, group := range config.Rooms {
 288                                 for _, room := range group {
 289                                         syncPowerLevels(client, room.ID, group, scoreData, counter%syncForceFrequency == 0)
 290                                 }
 291                         }
 292                         roomUsersMu.RUnlock()
 293                         counter++
 294                 }
 295         }()
 296         return client.Sync()
 297 }
 298
 299 func main() {
 300         err := Run()
 301         if err != nil {
 302                 log.Fatalf("Program failed: %v", err)
 303         }
 304 }