diff --git a/.gitignore b/.gitignore index 651d177..a456c74 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ config.json +stats.db clessy-broker clessy-mods clessy-stats \ No newline at end of file diff --git a/stats/main.go b/stats/main.go index a6c2610..cb3df54 100644 --- a/stats/main.go +++ b/stats/main.go @@ -14,20 +14,29 @@ func assert(err error) { } var db *bolt.DB +var chatID *int func process(broker *tg.Broker, update tg.APIMessage) { - + // Process messages from marked chat only + if update.Chat.ChatID != *chatID { + return + } + getNick(update.User) + updateStats(update) } func main() { brokerAddr := flag.String("broker", "localhost:7314", "Broker address:port") boltdbFile := flag.String("boltdb", "stats.db", "BoltDB database file") + chatID = flag.Int("chatid", -14625256, "Telegram Chat ID to count stats for") flag.Parse() - db, err := bolt.Open(*boltdbFile, 0600, nil) + var err error + db, err = bolt.Open(*boltdbFile, 0600, nil) assert(err) defer db.Close() + loadUsers() loadStats() err = tg.CreateBrokerClient(*brokerAddr, process) diff --git a/stats/stats.go b/stats/stats.go index b03f0a6..4ebe75d 100644 --- a/stats/stats.go +++ b/stats/stats.go @@ -7,6 +7,7 @@ import ( "time" "github.com/boltdb/bolt" + "github.com/hamcha/clessy/tg" ) const ( @@ -18,18 +19,21 @@ const ( MessageTypeVoice int = 5 MessageTypeContact int = 6 MessageTypeLocation int = 7 - MessageTypeMax int = 8 + MessageTypeDocument int = 8 + MessageTypeMax int = 9 ) type Stats struct { ByUserCount map[string]uint64 ByUserAvgLen map[string]uint64 + ByUserAvgCount map[string]uint64 ByWeekday [7]uint64 ByHour [24]uint64 ByType [MessageTypeMax]uint64 TodayDate time.Time Today uint64 TotalCount uint64 + TotalTxtCount uint64 TotalAvgLength uint64 Replies uint64 Forward uint64 @@ -52,30 +56,38 @@ func MakeUint(bval []byte, bucketName string, key string) uint64 { } } +func PutUint(value uint64) []byte { + bytes := make([]byte, 10) + n := binary.PutUvarint(bytes, value) + return bytes[:n] +} + func loadStats() { // Load today stats.TodayDate = time.Now() - err := db.View(func(tx *bolt.Tx) error { + err := db.Update(func(tx *bolt.Tx) error { b, err := tx.CreateBucketIfNotExists([]byte("global")) if err != nil { return err } // Load total messages counter - bval := bucket.Get([]byte("count")) + bval := b.Get([]byte("count")) stats.TotalCount = MakeUint(bval, "global", "count") - // Load total messages counter - bval := bucket.Get([]byte("avg")) + bval = b.Get([]byte("avg")) stats.TotalAvgLength = MakeUint(bval, "global", "avg") + bval = b.Get([]byte("avgcount")) + stats.TotalTxtCount = MakeUint(bval, "global", "avgcount") + // Load total replies counter - bval = bucket.Get([]byte("replies")) + bval = b.Get([]byte("replies")) stats.Replies = MakeUint(bval, "global", "replies") // Load total replies counter - bval = bucket.Get([]byte("forward")) + bval = b.Get([]byte("forward")) stats.Forward = MakeUint(bval, "global", "forward") // Load hour counters @@ -85,7 +97,7 @@ func loadStats() { } for i := 0; i < 24; i++ { - bval = bucket.Get([]byte(i)) + bval = b.Get([]byte{byte(i)}) stats.ByHour[i] = MakeUint(bval, "hour", strconv.Itoa(i)) } @@ -96,7 +108,7 @@ func loadStats() { } for i := 0; i < 7; i++ { - bval = bucket.Get([]byte(i)) + bval = b.Get([]byte{byte(i)}) stats.ByWeekday[i] = MakeUint(bval, "weekday", strconv.Itoa(i)) } @@ -107,7 +119,7 @@ func loadStats() { } todayKey := stats.TodayDate.Format("2006-1-2") - bval = bucket.Get([]byte(todayKey)) + bval = b.Get([]byte(todayKey)) stats.Today = MakeUint(bval, "date", todayKey) // Load user counters @@ -118,6 +130,7 @@ func loadStats() { } b.ForEach(func(user, messages []byte) error { stats.ByUserCount[string(user)] = MakeUint(messages, "users-count", string(user)) + return nil }) stats.ByUserAvgLen = make(map[string]uint64) @@ -127,6 +140,17 @@ func loadStats() { } b.ForEach(func(user, messages []byte) error { stats.ByUserAvgLen[string(user)] = MakeUint(messages, "users-avg", string(user)) + return nil + }) + + stats.ByUserAvgCount = make(map[string]uint64) + b, err = tx.CreateBucketIfNotExists([]byte("users-avgcount")) + if err != nil { + return err + } + b.ForEach(func(user, messages []byte) error { + stats.ByUserAvgCount[string(user)] = MakeUint(messages, "users-avgcount", string(user)) + return nil }) // Load type counters @@ -135,7 +159,7 @@ func loadStats() { return err } for i := 0; i < MessageTypeMax; i++ { - bval = bucket.Get([]byte(i)) + bval = b.Get([]byte{byte(i)}) stats.ByType[i] = MakeUint(bval, "types", strconv.Itoa(i)) } @@ -143,3 +167,225 @@ func loadStats() { }) assert(err) } + +func updateDate() { + err := db.Update(func(tx *bolt.Tx) error { + b := tx.Bucket([]byte("date")) + todayKey := stats.TodayDate.Format("2006-1-2") + err := b.Put([]byte(todayKey), PutUint(stats.Today)) + if err != nil { + return err + } + return nil + }) + if err != nil { + log.Println("[updateDate] Couldn't save last day stats: " + err.Error()) + } + stats.TodayDate = time.Now() + stats.Today = 0 +} + +func updateMean(currentMean, meanCount, newValue uint64) uint64 { + return ((currentMean * meanCount) + newValue) / (meanCount + 1) +} + +func updateStats(message tg.APIMessage) { + // + // Local update + // + + // DB Update flags + updatemean := false + updatetype := 0 + updatereplies := false + updateforward := false + + // Update total count + stats.TotalCount++ + + // Update individual user's count + username := message.User.Username + val, exists := stats.ByUserCount[username] + if !exists { + val = 0 + } + stats.ByUserCount[username] = val + 1 + + // Update time counters + now := time.Now() + hour := now.Hour() + wday := now.Weekday() + stats.ByHour[hour]++ + stats.ByWeekday[wday]++ + + // Check for day reset + if now.Day() != stats.TodayDate.Day() { + updateDate() + } + stats.Today++ + + // Text message + if message.Text != nil { + stats.ByType[MessageTypeText]++ + + // Update total and individual average + msglen := uint64(len(*(message.Text))) + if stats.TotalTxtCount > 0 { + stats.TotalAvgLength = updateMean(stats.TotalAvgLength, stats.TotalTxtCount, msglen) + stats.TotalTxtCount++ + } else { + stats.TotalAvgLength = msglen + stats.TotalTxtCount = 1 + } + val, exists = stats.ByUserAvgCount[username] + if exists { + stats.ByUserAvgLen[username] = updateMean(stats.ByUserAvgLen[username], val, msglen) + stats.ByUserAvgCount[username]++ + } else { + stats.ByUserAvgLen[username] = msglen + stats.ByUserAvgCount[username] = 1 + } + updatemean = true + updatetype = MessageTypeText + } + // Audio message + if message.Audio != nil { + stats.ByType[MessageTypeAudio]++ + updatetype = MessageTypeAudio + } + // Photo + if message.Photo != nil { + stats.ByType[MessageTypePhoto]++ + updatetype = MessageTypePhoto + } + // Sticker + if message.Sticker != nil { + stats.ByType[MessageTypeSticker]++ + updatetype = MessageTypeSticker + } + // Video + if message.Video != nil { + stats.ByType[MessageTypeVideo]++ + updatetype = MessageTypeVideo + } + // Voice message + if message.Voice != nil { + stats.ByType[MessageTypeVoice]++ + updatetype = MessageTypeVoice + } + // Contact + if message.Contact != nil { + stats.ByType[MessageTypeContact]++ + updatetype = MessageTypeContact + } + // Location + if message.Location != nil { + stats.ByType[MessageTypeLocation]++ + updatetype = MessageTypeLocation + } + // Document + if message.Document != nil { + stats.ByType[MessageTypeDocument]++ + updatetype = MessageTypeDocument + } + // Reply + if message.ReplyTo != nil { + stats.Replies++ + updatereplies = true + } + // Forwarded message + if message.FwdUser != nil { + stats.Forward++ + updateforward = true + } + + // + // DB Update + // + + err := db.Update(func(tx *bolt.Tx) error { + // Update total counters + b := tx.Bucket([]byte("global")) + + err := b.Put([]byte("count"), PutUint(stats.TotalCount)) + if err != nil { + return err + } + + if updatemean { + err = b.Put([]byte("avg"), PutUint(stats.TotalAvgLength)) + if err != nil { + return err + } + err = b.Put([]byte("avgcount"), PutUint(stats.TotalTxtCount)) + if err != nil { + return err + } + } + + if updatereplies { + err = b.Put([]byte("replies"), PutUint(stats.Replies)) + if err != nil { + return err + } + } + if updateforward { + err = b.Put([]byte("forward"), PutUint(stats.Forward)) + if err != nil { + return err + } + } + + // Update time counters + b = tx.Bucket([]byte("hour")) + err = b.Put([]byte{byte(hour)}, PutUint(stats.ByHour[hour])) + if err != nil { + return err + } + + b = tx.Bucket([]byte("weekday")) + err = b.Put([]byte{byte(wday)}, PutUint(stats.ByHour[wday])) + if err != nil { + return err + } + + b = tx.Bucket([]byte("date")) + todayKey := stats.TodayDate.Format("2006-1-2") + err = b.Put([]byte(todayKey), PutUint(stats.Today)) + if err != nil { + return err + } + + // Update user counters + b = tx.Bucket([]byte("users-count")) + err = b.Put([]byte(username), PutUint(stats.ByUserCount[username])) + if err != nil { + return err + } + + if updatemean { + b = tx.Bucket([]byte("users-avg")) + err = b.Put([]byte(username), PutUint(stats.ByUserAvgLen[username])) + if err != nil { + return err + } + b = tx.Bucket([]byte("users-avgcount")) + err = b.Put([]byte(username), PutUint(stats.ByUserAvgCount[username])) + if err != nil { + return err + } + } + + // Update type counter + b = tx.Bucket([]byte("types")) + err = b.Put([]byte{byte(updatetype)}, PutUint(stats.ByType[updatetype])) + if err != nil { + return err + } + + return nil + }) + if err != nil { + log.Println("[updateStats] Got error while updating DB: " + err.Error()) + } +} diff --git a/stats/users.go b/stats/users.go new file mode 100644 index 0000000..e160412 --- /dev/null +++ b/stats/users.go @@ -0,0 +1,47 @@ +package main + +import ( + "log" + "strings" + + "github.com/boltdb/bolt" + "github.com/hamcha/clessy/tg" +) + +var users map[string]string + +func getNick(apiuser tg.APIUser) { + if _, ok := users[apiuser.Username]; ok && strings.HasPrefix(users[apiuser.Username], apiuser.FirstName) { + // It's updated, don't bother + return + } + + users[apiuser.Username] = apiuser.FirstName + if apiuser.LastName != "" { + users[apiuser.Username] += apiuser.LastName + } + + err := db.Update(func(tx *bolt.Tx) error { + b := tx.Bucket([]byte("usernames")) + return b.Put([]byte(apiuser.Username), []byte(users[apiuser.Username])) + }) + if err != nil { + log.Printf("[getNick] Could not update %s name: %s\n", apiuser.Username, err.Error()) + } +} + +func loadUsers() { + users = make(map[string]string) + err := db.Update(func(tx *bolt.Tx) error { + b, err := tx.CreateBucketIfNotExists([]byte("usernames")) + if err != nil { + return err + } + b.ForEach(func(user, name []byte) error { + users[string(user)] = string(name) + return nil + }) + return nil + }) + assert(err) +}