Revamp TikTok support (Appears to now fully work)

AnnsAnns · Mar 15, 2024 · d3c409d · d3c409d
1 parent 80b7064
commit d3c409d
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 16 deletions.
diff --git a/sphene/src/main.rs b/sphene/src/main.rs
@@ -35,7 +35,7 @@ const REGEX_URL_EXTRACTOR: &str = r"\b(?:https?:\/\/|<)[^\s>]+(?:>|)\b";
 #[async_trait]
 impl EventHandler for Handler {
     async fn message(&self, context: Context, msg: Message) {
-        let mut url: String;
+        let url: String;
         let content = msg.content.clone();
         let options: Vec<CreateSelectMenuOption>;
 
@@ -61,7 +61,7 @@ impl EventHandler for Handler {
         } else if tiktok::is_tiktok_url(content.as_str())
             && self.dbconn.lock().await.get_server(id, false).tiktok
         {
-            url = tiktok::convert_url_lazy(content, tiktok::UrlType::VXTikTok).await;
+            url = tiktok::convert_url_lazy(tiktok::clear_url(content).await, tiktok::UrlType::TIKTXK).await;
             options = self.options_tiktok.clone();
         } else if instagram::is_instagram_url(content.as_str())
             && self.dbconn.lock().await.get_server(id, false).instagram
@@ -200,7 +200,7 @@ impl EventHandler for Handler {
                     .await
                 } else if tiktok::UrlType::from_string(&extracted_url) != tiktok::UrlType::Unknown {
                     tiktok::get_media_from_url(
-                        tiktok::convert_url_lazy(extracted_url, tiktok::UrlType::VXTikTok).await,
+                        tiktok::convert_url_lazy(extracted_url, tiktok::UrlType::TIKTXK).await,
                     )
                     .await
                 } else {
@@ -309,9 +309,9 @@ impl EventHandler for Handler {
                     new_msg,
                     bluesky::get_media_from_url(new_msg.clone()).await
                 );
-            } else if command == "direct_vxtiktok" {
+            } else if command == "direct_tiktxk" {
                 new_msg =
-                    tiktok::convert_url_lazy(extracted_url.to_string(), tiktok::UrlType::VXTikTok)
+                    tiktok::convert_url_lazy(extracted_url.to_string(), tiktok::UrlType::TIKTXK)
                         .await;
                 new_msg = format!(
                     "<{}> ({})",
@@ -404,8 +404,9 @@ async fn main() {
 
     let tiktok_options: Vec<CreateSelectMenuOption> = vec![
         download_option.clone(),
-        CreateSelectMenuOption::new("🔄️ Change to: VXTikTok", tiktok::VXTIKTOK_URL).to_owned(),
-        CreateSelectMenuOption::new("🖼️ Media Only", "direct_vxtiktok").to_owned(),
+        CreateSelectMenuOption::new("🔄️ Change to: TIKTXK", tiktok::TIKTXK_URL).to_owned(),
+        CreateSelectMenuOption::new("🔄️ Change to: TNKTOK", tiktok::TNKTOK_URL).to_owned(),
+        CreateSelectMenuOption::new("🖼️ Media Only", "direct_tiktxk").to_owned(),
         CreateSelectMenuOption::new("👶 Show original TikTok URL", tiktok::TIKTOK_URL).to_owned(),
         remove_option.clone(),
         disable_option.clone(),

diff --git a/thorium/src/tiktok.rs b/thorium/src/tiktok.rs
@@ -1,28 +1,34 @@
 pub const TIKTOK_URL: &str = "https://www.tiktok.com/";
+pub const TIKTOK_URL_TK: &str = "https://vt.tiktok.com/";
 pub const TIKTOK_NWWW_URL: &str = "https://tiktok.com/";
-pub const VXTIKTOK_URL: &str = "https://vxtiktok.com/";
+pub const TNKTOK_URL: &str = "https://tnktok.com/";
+pub const TIKTXK_URL: &str = "https://tiktxk.com/";
 
 #[derive(Debug, PartialEq)]
 pub enum UrlType {
     TikTok,
-    VXTikTok,
+    TNKTOK,
+    TIKTXK,
     Unknown
 }
 
 impl UrlType {
     pub fn as_str(&self) -> &'static str {
         match *self {
             UrlType::TikTok => TIKTOK_URL,
-            UrlType::VXTikTok => VXTIKTOK_URL,
+            UrlType::TIKTXK => TIKTXK_URL,
+            UrlType::TNKTOK => TNKTOK_URL,
             UrlType::Unknown => "",
         }
     }
 
     pub fn from_string(url: &str) -> UrlType {
-        if url.contains(TIKTOK_URL) || url.contains(TIKTOK_NWWW_URL) {
+        if url.contains(TIKTOK_URL) || url.contains(TIKTOK_NWWW_URL) || url.contains(TIKTOK_URL_TK) {
             return UrlType::TikTok;
-        } else if url.contains(VXTIKTOK_URL) {
-            return UrlType::VXTikTok;
+        } else if url.contains(TNKTOK_URL) {
+            return UrlType::TNKTOK;
+        } else if url.contains(TIKTXK_URL) {
+            return UrlType::TIKTXK;
         }
         UrlType::Unknown
     }
@@ -33,19 +39,34 @@ pub fn is_tiktok_url(url: &str) -> bool {
 }
 
 pub async fn convert_url(url: String, from: UrlType, to: UrlType) -> String {
-    url.replace(from.as_str(), to.as_str())
+    return clear_url(url).await.replace(from.as_str(), to.as_str());
 }
 
 pub async fn convert_url_lazy(url: String, to: UrlType) -> String {
-    let from = UrlType::from_string(&url);
+    let cleared_url = clear_url(url.clone()).await;
+    let from = UrlType::from_string(&cleared_url);
     url.replace(from.as_str(), to.as_str())
 }
 
+/**
+ * Remove vt. from the URL to get the original URL
+ */
+pub async fn clear_url(url: String) -> String {
+    if url.contains("vt.") {
+        // Remove vt. from the URL
+        return url.replace("vt.", "");
+    } else if !url.contains("www.") {
+        // Append www. to the URL
+        return url.replace("tiktok.com", "www.tiktok.com");
+    }
+    url
+}
+
 pub async fn get_url_type(url: String) -> UrlType {
     UrlType::from_string(&url)
 }
 
 pub async fn get_media_from_url(mut url: String) -> String {
-    url = convert_url_lazy(url, UrlType::VXTikTok).await;
+    url = convert_url_lazy(url, UrlType::TIKTXK).await;
     crate::get_media(url).await
 }