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