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 for %v handled from %v to %v - need restart", room.Name, evt.RoomID, tomb.ReplacementRoom)
 217                                 }
 218                         }
 219                 }
 220                 log.Printf("Room not found in config, so not doing room upgrade: %v", evt)
 221         })
 222         syncer.OnEventType(event.EventMessage, func(source mautrix.EventSource, evt *event.Event) {
 223                 if !isRoom(evt.RoomID) {
 224                         return
 225                 }
 226                 handleMessage(eventTime(evt), evt.RoomID, evt.Sender, evt)
 227         })
 228         syncer.OnEventType(event.StateMember, func(source mautrix.EventSource, evt *event.Event) {
 229                 if !isRoom(evt.RoomID) {
 230                         return
 231                 }
 232                 mem := evt.Content.AsMember()
 233                 key := evt.StateKey
 234                 if key == nil {
 235                         return
 236                 }
 237                 member := id.UserID(*key)
 238                 switch mem.Membership {
 239                 case event.MembershipJoin:
 240                         handleJoin(eventTime(evt), evt.RoomID, member, evt)
 241                 case event.MembershipLeave, event.MembershipBan:
 242                         handleLeave(eventTime(evt), evt.RoomID, member, evt)
 243                 default: // Ignore.
 244                 }
 245         })
 246         syncer.OnEventType(event.StatePowerLevels, func(source mautrix.EventSource, evt *event.Event) {
 247                 if !isRoom(evt.RoomID) {
 248                         return
 249                 }
 250                 handlePowerLevels(eventTime(evt), evt.RoomID, evt.Content.AsPowerLevels(), evt)
 251         })
 252         syncer.OnSync(func(resp *mautrix.RespSync, since string) bool {
 253                 // j, _ := json.MarshalIndent(resp, "", "  ")
 254                 // log.Print(string(j))
 255                 roomUsersMu.Lock()
 256                 if since != "" && !fullySynced {
 257                         log.Print("Fully synced.")
 258                         for room, users := range roomUsers {
 259                                 if _, found := users[config.UserID]; !found {
 260                                         log.Printf("Not actually joined %v yet...", room)
 261                                         _, err := client.JoinRoom(string(room), "", nil)
 262                                         if err != nil {
 263                                                 log.Printf("Failed to join %v: %v", room, err)
 264                                         }
 265                                 }
 266                         }
 267                         fullySynced = true
 268                 }
 269                 roomUsersMu.Unlock()
 270                 return true
 271         })
 272         client.Syncer = syncer
 273         ticker := time.NewTicker(syncInterval)
 274         defer ticker.Stop()
 275         go func() {
 276                 counter := 0
 277                 for range ticker.C {
 278                         roomUsersMu.RLock()
 279                         scoreData := map[id.RoomID]map[id.UserID]*Score{}
 280                         now := time.Now()
 281                         for room := range roomUsers {
 282                                 scores, err := queryUserScores(room, now)
 283                                 if err != nil {
 284                                         log.Fatalf("failed to query user scores: %v", err)
 285                                 }
 286                                 scoreData[room] = scores
 287                         }
 288                         for _, group := range config.Rooms {
 289                                 for _, room := range group {
 290                                         syncPowerLevels(client, room.ID, group, scoreData, counter%syncForceFrequency == 0)
 291                                 }
 292                         }
 293                         roomUsersMu.RUnlock()
 294                         counter++
 295                 }
 296         }()
 297         return client.Sync()
 298 }
 299
 300 func main() {
 301         err := Run()
 302         if err != nil {
 303                 log.Fatalf("Program failed: %v", err)
 304         }
 305 }