sauce_api/source/
iqdb.rs

1use std::time::Duration;
2
3use async_trait::async_trait;
4use reqwest::{header, Client};
5use visdom::{
6    types::{BoxDynElement, Elements},
7    Vis,
8};
9
10use crate::error::Error;
11
12use super::{Item, Output, Source};
13
14/// The [`IQDB`] source.
15///
16/// Works with `iqdb.org`
17#[derive(Debug)]
18pub struct Iqdb;
19
20#[async_trait]
21impl Source for Iqdb {
22    type Argument = ();
23
24    async fn check(&self, url: &str) -> Result<Output, Error> {
25        let client = Client::new();
26
27        // Check whether we're dealing with an image
28        let head = client.head(url).send().await?;
29
30        let content_type = head.headers().get(header::CONTENT_TYPE);
31
32        if let Some(content_type) = content_type {
33            let content_type = content_type.to_str()?;
34
35            if !content_type.contains("image") {
36                return Err(Error::LinkIsNotImage);
37            }
38        } else {
39            return Err(Error::LinkIsNotImage);
40        }
41
42        // Build the request
43
44        let req = client
45            .get("https://iqdb.org/")
46            .query(&[("url", url)])
47            .timeout(Duration::from_secs(10));
48
49        let resp = req.send().await?;
50
51        let text = resp.text().await?;
52
53        let html = Vis::load(text)?;
54
55        let pages = html.find("#pages").children("div");
56
57        let best_match = if pages.length() > 2 {
58            Self::harvest_best_match(&pages.eq(1))
59        } else {
60            None
61        };
62
63        let mut items = Vec::new();
64
65        if let Some(best_match) = best_match {
66            items.push(best_match);
67        }
68
69        for page in pages.into_iter().skip(2) {
70            let page = Self::harvest_page(&page);
71
72            items.extend(page);
73        }
74
75        for item in &mut items {
76            if item.link.starts_with("//") {
77                item.link = format!("https:{}", item.link);
78            }
79        }
80
81        Ok(Output {
82            original_url: url.to_string(),
83            items,
84        })
85    }
86
87    async fn create(_: Self::Argument) -> Result<Self, Error> {
88        Ok(Self)
89    }
90}
91
92impl Iqdb {
93    fn harvest_page(page: &BoxDynElement) -> Option<Item> {
94        let dom = Vis::dom(page);
95
96        let link = dom.find(".image a").first();
97
98        let url = link.attr("href")?;
99
100        let score = dom.find("tr");
101
102        if score.length() != 5 {
103            return Some(Item {
104                link: url.to_string(),
105                similarity: -1.0,
106            });
107        }
108
109        let score = score.eq(4);
110        let td = score.find("td");
111
112        let score = td.html();
113        let score = score.split_once('%')?.0.parse::<f32>().ok()? / 100.0;
114
115        Some(Item {
116            link: url.to_string(),
117            similarity: score,
118        })
119    }
120
121    fn harvest_best_match(pages: &Elements) -> Option<Item> {
122        let link = pages.find(".image a").first();
123
124        let url = link.attr("href")?;
125
126        let score = pages.find("tr");
127
128        if score.length() != 5 {
129            return Some(Item {
130                link: url.to_string(),
131                similarity: -1.0,
132            });
133        }
134
135        let score = score.eq(4);
136        let td = score.find("td");
137
138        let score = td.html();
139        let score = score.split_once('%')?.0.parse::<f32>().ok()? / 100.0;
140
141        Some(Item {
142            link: url.to_string(),
143            similarity: score,
144        })
145    }
146}
sauce_api/source/iqdb.rs

sauce_api/source/
iqdb.rs
