Inital commit

.gitignore

@@ -0,0 +1,5 @@
+/target
+config.toml
+.sqlx
+.idea
+data.db

Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "esix-database"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+anyhow = "1.0"
+clap = { version = "4.5.1", features = ["derive"] }
+config = "0.14"
+r621 = "0.2.3"
+serde = "1.0"
+serde_json = "1.0"
+sqlx = { version = "0.7.3", features = ["runtime-tokio-rustls", "sqlite"] }
+thiserror = "1.0"
+tokio = { version = "1.36", features = ["full"] }
+url = "2.5.0"

config.toml.example

@@ -0,0 +1,6 @@
+[database]
+url="sqlite://data.db"
+
+[e621]
+username=""
+apikey=""

migrations/20240225004801_init.sql

@@ -0,0 +1,14 @@
+CREATE TABLE IF NOT EXISTS chats (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL,
+    username TEXT
+) STRICT;
+
+CREATE TABLE IF NOT EXISTS posts (
+    id INTEGER PRIMARY KEY,
+    chat INTEGER NOT NULL,
+    chat_message_id INTEGER NOT NULL,
+    e621_id INTEGER NOT NULL,
+    FOREIGN KEY(chat) REFERENCES chats(id),
+    UNIQUE (chat, chat_message_id, e621_id)
+) STRICT;

src/config.rs

@@ -0,0 +1,45 @@
+use clap::{Parser, Subcommand};
+use config::{Config, File};
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct AppConfig {
+    pub database: DatabaseConfig,
+    pub e621: e621Config
+}
+
+impl AppConfig {
+    pub fn new() -> Self {
+        let config = Config::builder()
+            .add_source(File::with_name("config"))
+            .build()
+            .expect("Could not build config");
+
+        return config
+            .try_deserialize()
+            .expect("Could not deserialize config");
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct DatabaseConfig {
+    pub url: String,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct e621Config {
+    pub username: String,
+    pub apikey: String
+}
+
+#[derive(Debug, Parser)]
+pub struct Args {
+    #[command(subcommand)]
+    pub command: Command,
+}
+
+#[derive(Debug, Subcommand)]
+pub enum Command {
+    Import { path: String },
+    Unposted
+}

src/links.rs

@@ -0,0 +1,68 @@
+use std::str::FromStr;
+use thiserror::Error;
+use url::Url;
+
+#[derive(Error, Debug)]
+pub enum ParseLinkError {
+    #[error("Could not parse valid link")]
+    InvalidLink,
+
+    #[error("Unexpected host")]
+    UnknownHostError,
+}
+
+#[derive(Debug)]
+pub enum Link {
+    E621Post { post: u32 },
+    E621Pool { pool: u32 },
+}
+
+impl FromStr for Link {
+    type Err = ParseLinkError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let url = Url::from_str(s).map_err(|_| return ParseLinkError::InvalidLink)?;
+
+        if let Some(domain) = url.domain() {
+            match domain {
+                "www.e621.net" | "e621.net" | "www.e926.net" | "e926.net" => {
+                    if let Some(mut path_segments) = url.path_segments() {
+                        let first_segment = path_segments.next();
+                        if first_segment.is_none() {
+                            return Err(ParseLinkError::InvalidLink);
+                        }
+                        let first_segment = first_segment.unwrap();
+
+                        let last_segment = path_segments.last();
+                        if last_segment.is_none() {
+                            return Err(ParseLinkError::InvalidLink);
+                        }
+                        let last_segment = last_segment.unwrap();
+
+                        match first_segment {
+                            "post" | "posts" => {
+                                if let Ok(post) = u32::from_str(last_segment) {
+                                    return Ok(Link::E621Post { post });
+                                }
+                            }
+                            "pools" => {
+                                if let Ok(pool) = u32::from_str(last_segment) {
+                                    return Ok(Link::E621Pool { pool });
+                                }
+                            }
+                            _ => {
+                                println!("e621 first segment: {first_segment}");
+                            }
+                        }
+                    }
+                    return Err(ParseLinkError::UnknownHostError);
+                }
+                _ => {
+                    return Err(ParseLinkError::UnknownHostError);
+                }
+            }
+        } else {
+            return Err(ParseLinkError::UnknownHostError);
+        }
+    }
+}

src/main.rs

@@ -0,0 +1,145 @@
+use crate::config::{AppConfig, Args, Command};
+use crate::links::Link;
+use crate::telegram_chat_export::{TelegramChatExport, TextEntity};
+use clap::Parser;
+use sqlx::SqlitePool;
+use std::fs::OpenOptions;
+use std::str::FromStr;
+use r621::client::{Authentication, Client};
+use r621::post::Post;
+
+mod config;
+mod links;
+mod telegram_chat_export;
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    let config = AppConfig::new();
+    let args = Args::parse();
+
+    match args.command {
+        Command::Import { ref path } => {
+            import(path, config).await?;
+        },
+        Command::Unposted => {
+            list_unposted_favs(&config).await?;
+        }
+    }
+
+    Ok(())
+}
+
+async fn import(path: &String, config: AppConfig) -> anyhow::Result<()> {
+    let mut db = SqlitePool::connect_lazy(&config.database.url)?;
+    let import_file = OpenOptions::new().write(false).read(true).open(path)?;
+
+    let telegram_chat_export: TelegramChatExport = serde_json::from_reader(import_file)?;
+
+    import_chat(&telegram_chat_export, &mut db).await?;
+    import_messages(&telegram_chat_export, &mut db).await?;
+
+    Ok(())
+}
+
+async fn import_chat(
+    telegram_chat_export: &TelegramChatExport,
+    db: &SqlitePool,
+) -> anyhow::Result<()> {
+    let transaction = db.begin().await?;
+    sqlx::query!(
+        "INSERT OR REPLACE INTO chats(id, name) VALUES (?, ?)",
+        telegram_chat_export.id,
+        telegram_chat_export.name
+    )
+    .execute(db)
+    .await?;
+
+    transaction.commit().await?;
+    Ok(())
+}
+
+async fn import_messages(
+    telegram_chat_export: &TelegramChatExport,
+    db: &SqlitePool,
+) -> anyhow::Result<()> {
+    let transaction = db.begin().await?;
+    for message in &telegram_chat_export.messages {
+        for entity in &message.text_entities {
+            match entity {
+                TextEntity::Link { ref text } => {
+                    if let Ok(link) = Link::from_str(text) {
+                        match link {
+                            Link::E621Post { post } => {
+                                println!("{link:?}");
+                                insert_post(telegram_chat_export.id, message.id, post, db).await;
+                            }
+                            _ => {}
+                        }
+                    }
+                }
+                TextEntity::TextLink { text: _, ref href } => {
+                    if let Ok(link) = Link::from_str(href) {
+                        println!("{link:?}");
+                    }
+                }
+                _ => {}
+            }
+        }
+    }
+
+    transaction.commit().await?;
+    Ok(())
+}
+
+async fn insert_post(chat_id: u32, message_id: u32, post_id: u32, db: &SqlitePool) {
+    let result = sqlx::query!(
+        "INSERT INTO posts(chat, chat_message_id, e621_id) VALUES (?, ?, ?)",
+        chat_id,
+        message_id,
+        post_id
+    )
+    .execute(db)
+    .await;
+
+    if let Err(err) = result {
+        eprintln!("{err:?}");
+        eprintln!("{chat_id}, {message_id}, {post_id}");
+    }
+}
+
+async fn list_unposted_favs(config: &AppConfig) -> anyhow::Result<()> {
+    let db = SqlitePool::connect_lazy(&config.database.url)?;
+    let mut client = Client::new(
+        Authentication::Authorized {
+            username: &config.e621.username,
+            apikey: &config.e621.apikey
+        },
+        "esix-database/0.1 (by pixelhunter on e621)"
+    )?;
+
+    let mut esix_fav_posts = Vec::new();
+    let mut postcount = 1;
+    let mut page = 1;
+
+    while postcount > 0 {
+        println!("Fetching favorites page: {page}, Count: {}", esix_fav_posts.len());
+        let current_posts = client.list_posts(None, Some(format!("fav:{}", &config.e621.username)), Some(page)).await?;
+        postcount = current_posts.len();
+        esix_fav_posts.extend(current_posts);
+        page = page+1;
+    }
+
+    for fav_post in esix_fav_posts {
+        let post_id_i64 = fav_post.id as i64;
+        let result = sqlx::query!(
+            "SELECT COUNT(*) as count FROM posts WHERE e621_id = ?",
+            post_id_i64
+        ).fetch_one(&db).await?;
+
+        if result.count == 0 {
+            println!("https://www.e621.net/posts/{post_id_i64}");
+        }
+    }
+
+    Ok(())
+}

src/telegram_chat_export.rs

@@ -0,0 +1,28 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct TelegramChatExport {
+    pub name: String,
+    pub id: u32,
+    pub messages: Vec<Message>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Message {
+    pub id: u32,
+    pub text_entities: Vec<TextEntity>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum TextEntity {
+    Link { text: String },
+    TextLink { text: String, href: String },
+    Plain { text: String },
+    Hashtag { text: String },
+    Mention { text: String },
+    Bold { text: String },
+    Code { text: String },
+    Italic { text: String },
+    BotCommand { text: String },
+}