From aa40b93401f3c60adff2d77297d7bc55d3945e70 Mon Sep 17 00:00:00 2001
From: aNNiMON <annimon119@gmail.com>
Date: Mon, 20 Mar 2023 19:09:25 +0200
Subject: [PATCH] Add support for aibooru.online

---
 database.py  | 20 +++++++++++---------
 picsorter.py | 15 +++++++++------
 2 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/database.py b/database.py
index dd51576..27778ac 100644
--- a/database.py
+++ b/database.py
@@ -11,27 +11,29 @@ class Database:
         conn = sqlite3.connect(self.db_name)
         c = conn.cursor()
         c.executescript("""
-            CREATE TABLE IF NOT EXISTS danbooru (
-                id INTEGER PRIMARY KEY NOT NULL UNIQUE,
+            CREATE TABLE IF NOT EXISTS images (
+                id INTEGER PRIMARY KEY NOT NULL,
+                provider TEXT NOT NULL,
                 tags TEXT NOT NULL,
-                created_at TIMESTAMP
-            );
+                created_at TIMESTAMP,
+                UNIQUE(id, provider) ON CONFLICT REPLACE
+            )
             """)
         conn.commit()
         conn.close()
 
-    def is_exists(self, _id) -> bool:
+    def is_exists(self, _id, provider) -> bool:
         conn = sqlite3.connect(self.db_name)
         c = conn.cursor()
-        c.execute("SELECT EXISTS(SELECT 1 FROM danbooru WHERE id=?)", (_id, ))
+        c.execute("SELECT EXISTS(SELECT 1 FROM images WHERE id=? AND provider=?)", (_id, provider))
         result = c.fetchone()[0]
         conn.close()
         return bool(result)
     
-    def add(self, _id, tags):
+    def add(self, _id, provider, tags):
         conn = sqlite3.connect(self.db_name)
         c = conn.cursor()
-        sql = 'INSERT INTO danbooru(id, tags, created_at) VALUES (?,?,?)'
-        c.execute(sql, (_id, tags, datetime.now()))
+        sql = 'INSERT INTO images(id, provider, tags, created_at) VALUES (?,?,?,?)'
+        c.execute(sql, (_id, provider, tags, datetime.now()))
         conn.commit()
         conn.close()
diff --git a/picsorter.py b/picsorter.py
index 2bb5642..2e72ff9 100644
--- a/picsorter.py
+++ b/picsorter.py
@@ -95,20 +95,23 @@ class PicSorter:
         return url
 
     def __process_url(self, url: str) -> bool:
-        m = re.search(r"(?:posts/)?(\d{3,})", url)
+        m = re.search(r"(https://((?:dan|ai)booru|yande).*?(?:post(?:s|/show)/)?(\d{3,}))", url)
         if not m:
             return False
-        post_id = int(m.group(1))
-        if self.db.is_exists(post_id):
-            logging.info("Skipping exists post %d", post_id)
+        provider = m.group(2)
+        post_id = int(m.group(3))
+        if provider not in ['danbooru', 'aibooru']:
+            return False
+        if self.db.is_exists(provider, post_id):
+            logging.info("Skipping exists post %s %d", provider, post_id)
             return False
 
-        meta_result = self.metadata.process("https://danbooru.donmai.us/posts/" + str(post_id))
+        meta_result = self.metadata.process(m.group(1))
         if meta_result is None:
             return False
         image_path, tags = meta_result
         to_path = self.library.move(image_path, tags)
-        self.db.add(post_id, tags.tags_string)
+        self.db.add(post_id, provider, tags.tags_string)
         self.__show_path(to_path)
         return True