feat: support tags filtering

Closes #55
Samarium150 · Dec 6, 2021 · a44483c · a44483c
1 parent 6b339d1
commit a44483c
Show file tree

Hide file tree

Showing 9 changed files with 105 additions and 15 deletions.
diff --git a/src/main/kotlin/CommandConfig.kt b/src/main/kotlin/CommandConfig.kt
@@ -12,7 +12,7 @@ import net.mamoe.mirai.console.data.value
  * @constructor Create a Command config instance <br> 实例化命令配置
  * @see net.mamoe.mirai.console.data.AutoSavePluginConfig
  */
-object CommandConfig: AutoSavePluginConfig("CommandConfig") {
+object CommandConfig : AutoSavePluginConfig("CommandConfig") {
 
     /**
      * Composite command secondary names

diff --git a/src/main/kotlin/ImageData.kt b/src/main/kotlin/ImageData.kt
@@ -70,9 +70,9 @@ data class ImageData(
     fun toReadable(): String {
         return (
             "标题: ${title}\n" +
-            "作者: $author (uid: ${uid})\n" +
-            "标签: ${tags}\n" +
-            "链接: https://pixiv.net/artworks/${pid}"
-        )
+                "作者: $author (uid: ${uid})\n" +
+                "标签: ${tags}\n" +
+                "链接: https://pixiv.net/artworks/${pid}"
+            )
     }
 }
diff --git a/src/main/kotlin/Lolicon.kt b/src/main/kotlin/Lolicon.kt
@@ -42,7 +42,7 @@ import java.net.URL
  * @constructor Create a CompositeCommand instance <br> 实例化命令
  * @see net.mamoe.mirai.console.command.CompositeCommand
  */
-object Lolicon: CompositeCommand(
+object Lolicon : CompositeCommand(
     Main, primaryName = "lolicon",
     secondaryNames = CommandConfig.lolicon
 ) {
@@ -99,8 +99,10 @@ object Lolicon: CompositeCommand(
         val (r18, recall, cooldown) = ExecutionConfig.create(subject)
         val tags = tagArgs.joinToString(" ")
         val body = if (tags.isNotEmpty())
-            RequestBody(r18, 1, listOf(), "", Utils.processTags(tags),
-                listOf(PluginConfig.size), PluginConfig.proxy)
+            RequestBody(
+                r18, 1, listOf(), "", Utils.processTags(tags),
+                listOf(PluginConfig.size), PluginConfig.proxy
+            )
         else RequestBody(r18, 1, listOf(), tags, listOf(), listOf(PluginConfig.size), PluginConfig.proxy)
         Main.logger.info(body.toString())
         val response: ResponseBody?
@@ -122,6 +124,10 @@ object Lolicon: CompositeCommand(
         }
         try {
             val imageData = response.data[0]
+            if (!Utils.areTagsAllowed(imageData.tags)) {
+                sendMessage(ReplyConfig.filteredTag)
+                return
+            }
             val url = imageData.urls[PluginConfig.size] ?: return
             val imgInfoReceipt =
                 if (PluginConfig.verbose || subject == null) sendMessage(imageData.toReadable())
@@ -156,13 +162,14 @@ object Lolicon: CompositeCommand(
                         GlobalScope.launch {
                             Timer.cooldown(subject, cooldown)
                             withContext(Dispatchers.Default) {
-                                Main.logger.info(imgReceipt.target.toString()+"命令已冷却")
+                                Main.logger.info(imgReceipt.target.toString() + "命令已冷却")
                             }
                         }
                     }
                 }
             } catch (e: Exception) {
                 Main.logger.error(e)
+                sendMessage(ReplyConfig.networkError)
             } finally {
                 @Suppress("BlockingMethodInNonBlockingContext")
                 stream?.close()
@@ -263,6 +270,7 @@ object Lolicon: CompositeCommand(
                     imageInfoMsgBuilder.add("\n")
                 } catch (e: Exception) {
                     Main.logger.error(e)
+                    sendMessage(ReplyConfig.networkError)
                 } finally {
                     @Suppress("BlockingMethodInNonBlockingContext")
                     stream?.close()
@@ -286,7 +294,7 @@ object Lolicon: CompositeCommand(
                     GlobalScope.launch {
                         Timer.cooldown(subject, cooldown)
                         withContext(Dispatchers.Default) {
-                            Main.logger.info(imgReceipt.target.toString()+"命令已冷却")
+                            Main.logger.info(imgReceipt.target.toString() + "命令已冷却")
                         }
                     }
                 }

diff --git a/src/main/kotlin/Main.kt b/src/main/kotlin/Main.kt
@@ -38,7 +38,7 @@ import java.net.Proxy
  * @constructor Create a KotlinPlugin instance <br> 实例化插件
  * @see net.mamoe.mirai.console.plugin.jvm.KotlinPlugin
  */
-object Main: KotlinPlugin(
+object Main : KotlinPlugin(
     JvmPluginDescription(
         id = "com.github.samarium150.mirai-console-lolicon",
         version = "4.2.0",

diff --git a/src/main/kotlin/PluginConfig.kt b/src/main/kotlin/PluginConfig.kt
@@ -125,4 +125,20 @@ object PluginConfig : AutoSavePluginConfig("Config") {
      */
     @ValueDescription("默认的冷却时间(单位：s)")
     var cooldown: Int by value(60)
+
+    /**
+     * Tag filter mode
+     * <br>
+     * 标签过滤模式
+     */
+    @ValueDescription("标签过滤模式: none/whitelist/blacklist")
+    val tagFilterMode: String by value("none")
+
+    /**
+     * Tag filter
+     * <br>
+     * 标签过滤器
+     */
+    @ValueDescription("标签过滤器，白名单模式只发送包含指定标签的图片，黑名单模式只发送不包含指定标签的图片。(仅Get命令，Adv命令中不生效)")
+    val tagFilter: MutableSet<String> by value()
 }
diff --git a/src/main/kotlin/PluginData.kt b/src/main/kotlin/PluginData.kt
@@ -28,7 +28,7 @@ import net.mamoe.mirai.console.data.value
  * @constructor Create a plugin config instance <br> 实例化插件数据
  * @see net.mamoe.mirai.console.data.AutoSavePluginData
  */
-object PluginData: AutoSavePluginData("Data") {
+object PluginData : AutoSavePluginData("Data") {
 
     /**
      * User set

diff --git a/src/main/kotlin/ProxyConfig.kt b/src/main/kotlin/ProxyConfig.kt
@@ -12,7 +12,7 @@ import net.mamoe.mirai.console.data.value
  * @constructor Create empty Proxy config
  * @see net.mamoe.mirai.console.data.AutoSavePluginConfig
  */
-object ProxyConfig: AutoSavePluginConfig("ProxyConfig") {
+object ProxyConfig : AutoSavePluginConfig("ProxyConfig") {
 
     /**
      * Type

diff --git a/src/main/kotlin/ReplyConfig.kt b/src/main/kotlin/ReplyConfig.kt
@@ -12,7 +12,8 @@ import net.mamoe.mirai.console.data.value
  * @constructor Create a Reply config instance <br> 实例化回复配置
  * @see net.mamoe.mirai.console.data.AutoSavePluginConfig
  */
-object ReplyConfig: AutoSavePluginConfig("ReplyConfig") {
+object ReplyConfig : AutoSavePluginConfig("ReplyConfig") {
+
 
     /**
      * API returns errors
@@ -38,6 +39,13 @@ object ReplyConfig: AutoSavePluginConfig("ReplyConfig") {
     @ValueDescription("指令未冷却")
     val inCooldown: String by value("你怎么冲得到处都是")
 
+    /**
+     * Tag is filtered out
+     * <br>
+     * 标签被过滤
+     */
+    val filteredTag: String by value("该图片标签已被过滤，请换个标签再试")
+
     /**
      * Network error
      * <br>

diff --git a/src/main/kotlin/Utils.kt b/src/main/kotlin/Utils.kt
@@ -122,6 +122,64 @@ object Utils {
         return result.toList()
     }
 
+    /**
+     * Is the given [tag] allowed
+     *
+     * @param tag
+     * @return
+     */
+    private fun isTagAllowed(tag: String): Boolean {
+        return when (PluginConfig.tagFilterMode) {
+            "none" -> true
+            "whitelist" -> {
+                for (filter in PluginConfig.tagFilter) {
+                    if (filter.toRegex(setOf(RegexOption.IGNORE_CASE)).matches(tag)) return true
+                }
+                false
+            }
+            "blacklist" -> {
+                for (filter in PluginConfig.tagFilter) {
+                    if (filter.toRegex(setOf(RegexOption.IGNORE_CASE)).matches(tag)) return false
+                }
+                true
+            }
+            else -> true
+        }
+    }
+
+    /**
+     * Are these given [tags] allowed
+     *
+     * @param tags
+     * @return
+     */
+    fun areTagsAllowed(tags: List<String>): Boolean {
+        return when (PluginConfig.tagFilterMode) {
+            "none" -> true
+            "whitelist" -> {
+                var flag = false
+                for (tag in tags) {
+                    if (isTagAllowed(tag)) {
+                        flag = true
+                        break
+                    }
+                }
+                flag
+            }
+            "blacklist" -> {
+                var flag = true
+                for (tag in tags) {
+                    if (!isTagAllowed(tag)) {
+                        flag = false
+                        break
+                    }
+                }
+                flag
+            }
+            else -> true
+        }
+    }
+
     private val sizeMap: Map<String, Int> = mapOf(
         "original" to 0,
         "regular" to 1,
@@ -131,7 +189,7 @@ object Utils {
     )
 
     fun getUrl(urls: Map<String, String>): String? {
-        return urls[urls.keys.sortedBy { sizeMap[it] } [0]]
+        return urls[urls.keys.sortedBy { sizeMap[it] }[0]]
     }
 
     /**