From e690edfccb398c18f78dc0b3c378509c9d114efd Mon Sep 17 00:00:00 2001 From: djaler Date: Mon, 29 Jan 2024 20:49:39 +0300 Subject: [PATCH] =?UTF-8?q?feat:=20=D0=BE=D1=82=D1=81=D0=BB=D0=B5=D0=B6?= =?UTF-8?q?=D0=B8=D0=B2=D0=B0=D0=BD=D0=B8=D0=B5=20=D0=B4=D1=83=D0=B1=D0=BB?= =?UTF-8?q?=D0=B8=D0=BA=D0=B0=D1=82=D0=BE=D0=B2=20=D0=B8=D0=B7=D0=BE=D0=B1?= =?UTF-8?q?=D1=80=D0=B0=D0=B6=D0=B5=D0=BD=D0=B8=D0=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../github/djaler/evilbot/entity/ImageHash.kt | 23 +++++++ .../evilbot/handlers/SeenMemeHandler.kt | 53 ++++++++++++++++ .../evilbot/repository/ImageHashRepository.kt | 8 +++ .../evilbot/service/DuplicateImageChecker.kt | 63 +++++++++++++++++++ .../migration/V9__add_image_hashes_table.sql | 10 +++ 5 files changed, 157 insertions(+) create mode 100644 src/main/kotlin/com/github/djaler/evilbot/entity/ImageHash.kt create mode 100644 src/main/kotlin/com/github/djaler/evilbot/handlers/SeenMemeHandler.kt create mode 100644 src/main/kotlin/com/github/djaler/evilbot/repository/ImageHashRepository.kt create mode 100644 src/main/kotlin/com/github/djaler/evilbot/service/DuplicateImageChecker.kt create mode 100644 src/main/resources/db/migration/V9__add_image_hashes_table.sql diff --git a/src/main/kotlin/com/github/djaler/evilbot/entity/ImageHash.kt b/src/main/kotlin/com/github/djaler/evilbot/entity/ImageHash.kt new file mode 100644 index 0000000..22bf1cd --- /dev/null +++ b/src/main/kotlin/com/github/djaler/evilbot/entity/ImageHash.kt @@ -0,0 +1,23 @@ +package com.github.djaler.evilbot.entity + +import javax.persistence.* + +@Entity +@Table(name = "image_hashes") +data class ImageHash( + @Column + val hash: String, + + @Column + val chatId: Short, + + @Column + val messageId: Long, + + @Column + val fileId: String, + + @Id + @GeneratedValue(strategy = GenerationType.IDENTITY) + val id: Int = 0 +) diff --git a/src/main/kotlin/com/github/djaler/evilbot/handlers/SeenMemeHandler.kt b/src/main/kotlin/com/github/djaler/evilbot/handlers/SeenMemeHandler.kt new file mode 100644 index 0000000..03f1828 --- /dev/null +++ b/src/main/kotlin/com/github/djaler/evilbot/handlers/SeenMemeHandler.kt @@ -0,0 +1,53 @@ +package com.github.djaler.evilbot.handlers + +import com.github.djaler.evilbot.handlers.base.MessageHandler +import com.github.djaler.evilbot.service.ChatService +import com.github.djaler.evilbot.service.DuplicateImageChecker +import dev.inmo.tgbotapi.bot.RequestsExecutor +import dev.inmo.tgbotapi.extensions.api.files.downloadFile +import dev.inmo.tgbotapi.extensions.api.send.reply +import dev.inmo.tgbotapi.extensions.utils.asContentMessage +import dev.inmo.tgbotapi.extensions.utils.asPhotoContent +import dev.inmo.tgbotapi.extensions.utils.asPublicChat +import dev.inmo.tgbotapi.extensions.utils.formatting.makeLinkToMessage +import dev.inmo.tgbotapi.types.message.abstracts.Message +import dev.inmo.tgbotapi.utils.buildEntities +import dev.inmo.tgbotapi.utils.link +import org.springframework.stereotype.Component +import java.io.ByteArrayInputStream +import javax.imageio.ImageIO + +@Component +class SeenMemeHandler( + private val requestExecutor: RequestsExecutor, + private val duplicateImageChecker: DuplicateImageChecker, + private val chatService: ChatService, +) : MessageHandler() { + override suspend fun handleMessage(message: Message): Boolean { + val chat = message.chat.asPublicChat() ?: return false + val imageFile = message.asContentMessage()?.content?.asPhotoContent()?.media ?: return false + + val photoBytes = requestExecutor.downloadFile(imageFile) + val image = ByteArrayInputStream(photoBytes).use { + ImageIO.read(it) + } + + val (chatEntity, _) = chatService.getOrCreateChatFrom(chat) + val originalMessageId: Long? = duplicateImageChecker.findDuplicate(image, chatEntity) + + if (originalMessageId == null) { + duplicateImageChecker.saveHash(image, chatEntity, message.messageId, imageFile.fileId) + return false + } else { + val messageLink = makeLinkToMessage(message.chat, originalMessageId) ?: return false + + requestExecutor.reply( + message, + buildEntities { + +"Уже было - " + link(messageLink) + } + ) + return true + } + } +} diff --git a/src/main/kotlin/com/github/djaler/evilbot/repository/ImageHashRepository.kt b/src/main/kotlin/com/github/djaler/evilbot/repository/ImageHashRepository.kt new file mode 100644 index 0000000..cd46f6f --- /dev/null +++ b/src/main/kotlin/com/github/djaler/evilbot/repository/ImageHashRepository.kt @@ -0,0 +1,8 @@ +package com.github.djaler.evilbot.repository + +import com.github.djaler.evilbot.entity.ImageHash +import org.springframework.data.jpa.repository.JpaRepository + +interface ImageHashRepository : JpaRepository { + fun findByChatIdAndHash(chatId: Short, hash: String): ImageHash? +} diff --git a/src/main/kotlin/com/github/djaler/evilbot/service/DuplicateImageChecker.kt b/src/main/kotlin/com/github/djaler/evilbot/service/DuplicateImageChecker.kt new file mode 100644 index 0000000..6d3a999 --- /dev/null +++ b/src/main/kotlin/com/github/djaler/evilbot/service/DuplicateImageChecker.kt @@ -0,0 +1,63 @@ +package com.github.djaler.evilbot.service + +import com.github.djaler.evilbot.entity.Chat +import com.github.djaler.evilbot.entity.ImageHash +import com.github.djaler.evilbot.repository.ImageHashRepository +import dev.inmo.tgbotapi.requests.abstracts.FileId +import dev.inmo.tgbotapi.types.MessageIdentifier +import korlibs.crypto.encoding.hex +import org.springframework.stereotype.Component +import java.awt.image.BufferedImage +import java.io.ByteArrayOutputStream +import java.security.DigestOutputStream +import java.security.MessageDigest +import javax.imageio.ImageIO + +@Component +class DuplicateImageChecker( + private val imageHashRepository: ImageHashRepository, +) { + fun findDuplicate(image: BufferedImage, chat: Chat): MessageIdentifier? { + val hash = resizeAndGetHash(image) + + val duplicate = imageHashRepository.findByChatIdAndHash(chat.id, hash) + + return duplicate?.messageId + } + + fun saveHash(image: BufferedImage, chat: Chat, messageId: MessageIdentifier, fileId: FileId) { + val hash = resizeAndGetHash(image) + + imageHashRepository.save(ImageHash( + hash, + chat.id, + messageId, + fileId.fileId + )) + } + + // TODO cache + private fun resizeAndGetHash(image: BufferedImage): String { + val resizedImage = resizeImage(image, width = 64, height = 64) + + return getImageHash(resizedImage) + } + + private fun resizeImage(image: BufferedImage, width: Int, height: Int): BufferedImage { + return BufferedImage(width, height, image.type).apply { + graphics.drawImage(image, 0, 0, 64, 64, null) + } + } + + private fun getImageHash(image: BufferedImage): String { + val messageDigest = MessageDigest.getInstance("MD5") + ByteArrayOutputStream().use { baos -> + DigestOutputStream(baos, messageDigest).use { dos -> + ImageIO.write(image, "png", dos) + dos.flush() + } + } + + return messageDigest.digest().hex + } +} diff --git a/src/main/resources/db/migration/V9__add_image_hashes_table.sql b/src/main/resources/db/migration/V9__add_image_hashes_table.sql new file mode 100644 index 0000000..b11f4ef --- /dev/null +++ b/src/main/resources/db/migration/V9__add_image_hashes_table.sql @@ -0,0 +1,10 @@ +CREATE TABLE image_hashes +( + id SERIAL PRIMARY KEY, + hash VARCHAR NOT NULL, + chat_id SMALLINT NOT NULL, + message_id BIGINT NOT NULL, + file_id VARCHAR NOT NULL, + + UNIQUE (chat_id, hash) +);