Telegram bot to find previous images from a chat
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

bot.go 5.5KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287
  1. package main
  2. import (
  3. "encoding/json"
  4. "fmt"
  5. "io"
  6. "net/http"
  7. "os"
  8. "strconv"
  9. "strings"
  10. "image"
  11. _ "image/jpeg"
  12. _ "image/png"
  13. "github.com/boltdb/bolt"
  14. "github.com/corona10/goimagehash"
  15. "github.com/Syfaro/finch"
  16. "github.com/go-telegram-bot-api/telegram-bot-api"
  17. _ "github.com/Syfaro/finch/commands/help"
  18. )
  19. var db *bolt.DB
  20. type replier interface {
  21. QuickReply(tgbotapi.Message, string) error
  22. }
  23. func processReader(msg tgbotapi.Message, cmd replier, r io.Reader) error {
  24. img, _, err := image.Decode(r)
  25. if err != nil {
  26. return nil // ignore bad images, some links might not be images
  27. }
  28. hash, err := goimagehash.PerceptionHash(img)
  29. if err != nil {
  30. return err
  31. }
  32. chat := []byte(strconv.FormatInt(msg.Chat.ID, 10))
  33. err = db.Update(func(tx *bolt.Tx) error {
  34. b, err := tx.CreateBucketIfNotExists(chat)
  35. if err != nil {
  36. return err
  37. }
  38. by, err := json.Marshal(msg)
  39. if err != nil {
  40. return err
  41. }
  42. return b.Put(by, []byte(hash.ToString()))
  43. })
  44. if err != nil {
  45. return err
  46. }
  47. err = db.View(func(tx *bolt.Tx) error {
  48. b := tx.Bucket(chat)
  49. thresholds := tx.Bucket([]byte("thresholds"))
  50. return b.ForEach(func(k, v []byte) error {
  51. h, err := goimagehash.ImageHashFromString(string(v))
  52. if err != nil {
  53. return err
  54. }
  55. var oldMessage tgbotapi.Message
  56. err = json.Unmarshal(k, &oldMessage)
  57. if err != nil {
  58. return err
  59. }
  60. if oldMessage.MessageID == msg.MessageID {
  61. return nil
  62. }
  63. diff, err := hash.Distance(h)
  64. if err != nil {
  65. return err
  66. }
  67. thres := thresholds.Get(chat)
  68. if thres == nil {
  69. return nil
  70. }
  71. val, err := strconv.Atoi(string(thres))
  72. if err != nil {
  73. return err
  74. }
  75. if diff > val {
  76. return nil
  77. }
  78. return cmd.QuickReply(oldMessage, fmt.Sprintf("Has similarity of %d", diff))
  79. })
  80. })
  81. if err != nil {
  82. return err
  83. }
  84. return nil
  85. }
  86. type handleMessage struct {
  87. finch.CommandBase
  88. }
  89. func (cmd handleMessage) ShouldExecute(tgbotapi.Message) bool {
  90. return true
  91. }
  92. func (cmd handleMessage) Execute(msg tgbotapi.Message) error {
  93. if msg.Photo != nil {
  94. pixels := 0
  95. idx := -1
  96. for i, p := range *msg.Photo { // find largest image in set
  97. pix := p.Height * p.Width
  98. if pix > pixels {
  99. pixels = pix
  100. idx = i
  101. }
  102. }
  103. if idx == -1 { // no photo sizes
  104. return nil
  105. }
  106. url, err := cmd.API.GetFileDirectURL((*msg.Photo)[idx].FileID)
  107. if err != nil {
  108. return err
  109. }
  110. resp, err := http.Get(url)
  111. if err != nil {
  112. return err
  113. }
  114. return processReader(msg, cmd, resp.Body)
  115. } else if msg.Document != nil {
  116. url, err := cmd.API.GetFileDirectURL(msg.Document.FileID)
  117. if err != nil {
  118. return err
  119. }
  120. resp, err := http.Get(url)
  121. if err != nil {
  122. return err
  123. }
  124. return processReader(msg, cmd, resp.Body)
  125. } else if entities := msg.Entities; entities != nil {
  126. for _, entity := range *entities {
  127. var link string
  128. if entity.Type == "text_link" {
  129. link = entity.URL
  130. } else if entity.Type == "url" {
  131. link = msg.Text[entity.Offset : entity.Offset+entity.Length]
  132. }
  133. if link == "" {
  134. continue
  135. }
  136. resp, err := http.Head(link)
  137. if err != nil {
  138. continue
  139. }
  140. switch resp.Header.Get("Content-Type") {
  141. case "image/jpeg", "image/png":
  142. resp, err := http.Get(link)
  143. if err != nil {
  144. return err
  145. }
  146. if err = processReader(msg, cmd, resp.Body); err != nil {
  147. return err
  148. }
  149. default:
  150. continue
  151. }
  152. }
  153. }
  154. return nil
  155. }
  156. func (handleMessage) Help() finch.Help {
  157. return finch.Help{
  158. Name: "Background Image Loader",
  159. }
  160. }
  161. type handleThreshold struct {
  162. finch.CommandBase
  163. }
  164. func (handleThreshold) ShouldExecute(msg tgbotapi.Message) bool {
  165. return finch.SimpleArgCommand("threshold", 1, msg.Text)
  166. }
  167. func (cmd handleThreshold) Execute(msg tgbotapi.Message) error {
  168. return db.Update(func(tx *bolt.Tx) error {
  169. b, err := tx.CreateBucketIfNotExists([]byte("thresholds"))
  170. if err != nil {
  171. return err
  172. }
  173. args := strings.Fields(msg.CommandArguments())
  174. if len(args) != 1 {
  175. return cmd.QuickReply(msg, "You must include the threshold value")
  176. }
  177. arg := args[0]
  178. if _, err := strconv.Atoi(arg); err != nil {
  179. return cmd.QuickReply(msg, "Your threshold must be a number")
  180. }
  181. if err = b.Put([]byte(strconv.FormatInt(msg.Chat.ID, 10)), []byte(arg)); err != nil {
  182. return err
  183. }
  184. return cmd.QuickReply(msg, "Updated threshold!")
  185. })
  186. }
  187. func (handleThreshold) Help() finch.Help {
  188. return finch.Help{
  189. Name: "Set Threshold",
  190. Description: "Set minimum threshold to bring up an old message",
  191. Botfather: [][]string{[]string{"threshold", "Set the minimum threshold for an old message"}},
  192. Example: "/threshold 15",
  193. }
  194. }
  195. type handleStart struct {
  196. finch.CommandBase
  197. }
  198. func (handleStart) ShouldExecute(msg tgbotapi.Message) bool {
  199. return finch.SimpleCommand("start", msg.Text)
  200. }
  201. func (cmd handleStart) Execute(msg tgbotapi.Message) error {
  202. return cmd.QuickReply(msg, "Hello! Add me to a group and set your similarity threshold to get started. 0 means identical, 10-15 is a good starting point.")
  203. }
  204. func (handleStart) Help() finch.Help {
  205. return finch.Help{
  206. Name: "Start",
  207. }
  208. }
  209. func main() {
  210. finch.RegisterCommand(&handleStart{})
  211. finch.RegisterCommand(&handleMessage{})
  212. finch.RegisterCommand(&handleThreshold{})
  213. d, err := bolt.Open(os.Getenv("BOLT_DBNAME"), 0600, nil)
  214. if err != nil {
  215. panic(err)
  216. }
  217. defer d.Close()
  218. db = d
  219. if err = db.Update(func(tx *bolt.Tx) error {
  220. _, err := tx.CreateBucketIfNotExists([]byte("thresholds"))
  221. return err
  222. }); err != nil {
  223. panic(err)
  224. }
  225. bot := finch.NewFinch(os.Getenv("TELEGRAM_APITOKEN"))
  226. bot.API.Debug = os.Getenv("TELEGRAM_DEBUG") == "true"
  227. bot.Start()
  228. }