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         // Note: we have to lower case the user ID for Matrix protocol communication.
  55         uid := id.UserID(strings.ToLower(string(config.UserID)))
  56         client, err := mautrix.NewClient(config.Homeserver, uid, config.AccessToken)
  57         if err != nil {
  58                 return nil, fmt.Errorf("failed to create client: %v", err)
  59         }
  60         if config.AccessToken == "" {
  61                 resp, err := client.Login(&mautrix.ReqLogin{
  62                         Type: mautrix.AuthTypePassword,
  63                         Identifier: mautrix.UserIdentifier{
  64                                 Type: mautrix.IdentifierTypeUser,
  65                                 User: string(client.UserID),
  66                         },
  67                         Password:                 config.Password,
  68                         InitialDeviceDisplayName: "matrixbot",
  69                         StoreCredentials:         true,
  70                 })
  71                 if err != nil {
  72                         return nil, fmt.Errorf("failed to authenticate: %v", err)
  73                 }
  74                 config.Password = ""
  75                 config.DeviceID = resp.DeviceID
  76                 config.AccessToken = resp.AccessToken
  77                 err = config.Save()
  78                 if err != nil {
  79                         return nil, fmt.Errorf("failed to save config: %v", err)
  80                 }
  81         } else {
  82                 client.DeviceID = config.DeviceID
  83         }
  84         return client, nil
  85 }
  86
  87 var (
  88         roomUsers       = map[id.RoomID]map[id.UserID]struct{}{}
  89         roomUsersMu     sync.RWMutex
  90         fullySynced     bool
  91         roomPowerLevels = map[id.RoomID]*event.PowerLevelsEventContent{}
  92 )
  93
  94 func setUserStateAt(room id.RoomID, user id.UserID, now time.Time, maxPrevState, state State) {
  95         err := writeUserStateAt(room, user, now, maxPrevState, state)
  96         if err != nil {
  97                 log.Fatalf("failed to write user state: %v", err)
  98         }
  99 }
 100
 101 func handleMessage(now time.Time, room id.RoomID, sender id.UserID, raw *event.Event) {
 102         // log.Printf("[%v] Message from %v to %v", now, sender, room)
 103         roomUsersMu.Lock()
 104         roomUsers[room][sender] = struct{}{}
 105         roomUsersMu.Unlock()
 106         setUserStateAt(room, sender, now.Add(-activeTime), Active, Active)
 107         setUserStateAt(room, sender, now, Active, Idle)
 108 }
 109
 110 func handleJoin(now time.Time, room id.RoomID, member id.UserID, raw *event.Event) {
 111         log.Printf("[%v] Join from %v to %v", now, member, room)
 112         roomUsersMu.Lock()
 113         roomUsers[room][member] = struct{}{}
 114         roomUsersMu.Unlock()
 115         setUserStateAt(room, member, now, NotActive, Idle)
 116 }
 117
 118 func handleLeave(now time.Time, room id.RoomID, member id.UserID, raw *event.Event) {
 119         log.Printf("[%v] Leave from %v to %v", now, member, room)
 120         roomUsersMu.Lock()
 121         delete(roomUsers[room], member)
 122         roomUsersMu.Unlock()
 123         setUserStateAt(room, member, now, Active, NotActive)
 124 }
 125
 126 func handlePowerLevels(now time.Time, room id.RoomID, levels *event.PowerLevelsEventContent, raw *event.Event) {
 127         // log.Printf("[%v] Power levels for %v are %v", now, room, levels)
 128         levelsCopy := *levels // Looks like mautrix always passes the same pointer here.
 129         roomUsersMu.Lock()
 130         roomPowerLevels[room] = &levelsCopy
 131         roomUsersMu.Unlock()
 132 }
 133
 134 func eventTime(evt *event.Event) time.Time {
 135         return time.Unix(0, evt.Timestamp*1000000)
 136 }
 137
 138 type MoreMessagesSyncer struct {
 139         *mautrix.DefaultSyncer
 140 }
 141
 142 func newSyncer() *MoreMessagesSyncer {
 143         return &MoreMessagesSyncer{
 144                 DefaultSyncer: mautrix.NewDefaultSyncer(),
 145         }
 146 }
 147
 148 func (s *MoreMessagesSyncer) GetFilterJSON(userID id.UserID) *mautrix.Filter {
 149         f := s.DefaultSyncer.GetFilterJSON(userID)
 150         // Same filters as Element.
 151         f.Room.Timeline.Limit = 20
 152         // Only include our rooms.
 153         f.Room.Rooms = make([]id.RoomID, 0, len(roomUsers))
 154         for room := range roomUsers {
 155                 f.Room.Rooms = append(f.Room.Rooms, room)
 156         }
 157         return f
 158 }
 159
 160 func isRoom(room id.RoomID) bool {
 161         roomUsersMu.RLock()
 162         defer roomUsersMu.RUnlock()
 163         _, found := roomUsers[room]
 164         return found
 165 }
 166
 167 func Run() (err error) {
 168         err = InitDatabase()
 169         if err != nil {
 170                 return fmt.Errorf("failed to init database: %v", err)
 171         }
 172         defer func() {
 173                 err2 := CloseDatabase()
 174                 if err2 != nil && err == nil {
 175                         err = fmt.Errorf("failed to close database: %v", err)
 176                 }
 177         }()
 178         logPowerLevelBounds()
 179         config := &Config{}
 180         err = config.Load()
 181         if err != nil {
 182                 return fmt.Errorf("failed to load config: %v", err)
 183         }
 184         for _, group := range config.Rooms {
 185                 for _, room := range group {
 186                         roomUsers[room.ID] = map[id.UserID]struct{}{}
 187                 }
 188         }
 189         client, err := Login(config)
 190         if err != nil {
 191                 return fmt.Errorf("failed to login: %v", err)
 192         }
 193         syncer := newSyncer()
 194         syncer.OnEventType(event.EventMessage, func(source mautrix.EventSource, evt *event.Event) {
 195                 if !isRoom(evt.RoomID) {
 196                         return
 197                 }
 198                 handleMessage(eventTime(evt), evt.RoomID, evt.Sender, evt)
 199         })
 200         syncer.OnEventType(event.StateMember, func(source mautrix.EventSource, evt *event.Event) {
 201                 if !isRoom(evt.RoomID) {
 202                         return
 203                 }
 204                 mem := evt.Content.AsMember()
 205                 key := evt.StateKey
 206                 if key == nil {
 207                         return
 208                 }
 209                 member := id.UserID(*key)
 210                 switch mem.Membership {
 211                 case event.MembershipJoin:
 212                         handleJoin(eventTime(evt), evt.RoomID, member, evt)
 213                 case event.MembershipLeave, event.MembershipBan:
 214                         handleLeave(eventTime(evt), evt.RoomID, member, evt)
 215                 default: // Ignore.
 216                 }
 217         })
 218         syncer.OnEventType(event.StatePowerLevels, func(source mautrix.EventSource, evt *event.Event) {
 219                 if !isRoom(evt.RoomID) {
 220                         return
 221                 }
 222                 handlePowerLevels(eventTime(evt), evt.RoomID, evt.Content.AsPowerLevels(), evt)
 223         })
 224         syncer.OnSync(func(resp *mautrix.RespSync, since string) bool {
 225                 // j, _ := json.MarshalIndent(resp, "", "  ")
 226                 // log.Print(string(j))
 227                 roomUsersMu.Lock()
 228                 if since != "" && !fullySynced {
 229                         log.Print("Fully synced.")
 230                         for room, users := range roomUsers {
 231                                 if _, found := users[config.UserID]; !found {
 232                                         log.Printf("Not actually joined %v yet...", room)
 233                                         _, err := client.JoinRoom(string(room), "", nil)
 234                                         if err != nil {
 235                                                 log.Printf("Failed to join %v: %v", room, err)
 236                                         }
 237                                 }
 238                         }
 239                         fullySynced = true
 240                 }
 241                 roomUsersMu.Unlock()
 242                 return true
 243         })
 244         client.Syncer = syncer
 245         ticker := time.NewTicker(syncInterval)
 246         defer ticker.Stop()
 247         go func() {
 248                 counter := 0
 249                 for range ticker.C {
 250                         roomUsersMu.RLock()
 251                         scoreData := map[id.RoomID]map[id.UserID]*Score{}
 252                         now := time.Now()
 253                         for room := range roomUsers {
 254                                 scores, err := queryUserScores(room, now)
 255                                 if err != nil {
 256                                         log.Fatalf("failed to query user scores: %v", err)
 257                                 }
 258                                 scoreData[room] = scores
 259                         }
 260                         for _, group := range config.Rooms {
 261                                 for _, room := range group {
 262                                         syncPowerLevels(client, room.ID, group, scoreData, counter%syncForceFrequency == 0)
 263                                 }
 264                         }
 265                         roomUsersMu.RUnlock()
 266                         counter++
 267                 }
 268         }()
 269         return client.Sync()
 270 }
 271
 272 func main() {
 273         err := Run()
 274         if err != nil {
 275                 log.Fatalf("Program failed: %v", err)
 276         }
 277 }