Skip to content

Commit

Permalink
feat: отслеживание дубликатов изображений
Browse files Browse the repository at this point in the history
  • Loading branch information
Djaler committed Jan 29, 2024
1 parent 3caa5da commit e690edf
Show file tree
Hide file tree
Showing 5 changed files with 157 additions and 0 deletions.
23 changes: 23 additions & 0 deletions src/main/kotlin/com/github/djaler/evilbot/entity/ImageHash.kt
Original file line number Diff line number Diff line change
@@ -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
)
Original file line number Diff line number Diff line change
@@ -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
}
}
}
Original file line number Diff line number Diff line change
@@ -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<ImageHash, Int> {
fun findByChatIdAndHash(chatId: Short, hash: String): ImageHash?
}
Original file line number Diff line number Diff line change
@@ -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
}
}
10 changes: 10 additions & 0 deletions src/main/resources/db/migration/V9__add_image_hashes_table.sql
Original file line number Diff line number Diff line change
@@ -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)
);

0 comments on commit e690edf

Please sign in to comment.