Skip to main content

scrapr_core/
protocols.rs

1use crate::request::*;
2use anyhow::{Result, anyhow};
3use native_tls::TlsConnector;
4use std::{
5    io::{Read, Write},
6    net::TcpStream,
7};
8
9/// Fetch a page over plain HTTP.
10///
11/// # Arguments
12///
13/// * `host` - The domain name of the server.
14/// * `path` - The path to fetch (e.g., "/index.html").
15///
16/// # Returns
17///
18/// Body of the HTTP response if successful.
19pub fn fetch_http(host: &str, path: &str) -> Result<String> {
20    let mut stream =
21        TcpStream::connect((host, 80)).map_err(|e| anyhow!("Failed to connect: {e}"))?;
22
23    let request = format!("GET {path} HTTP/1.1\r\nHost: {host}\r\nConnection: close\r\n\r\n");
24    stream.write_all(request.as_bytes())?;
25
26    let mut response = String::new();
27    stream.read_to_string(&mut response)?;
28
29    response
30        .split("\r\n\r\n")
31        .nth(1)
32        .map(|body| body.to_string())
33        .ok_or_else(|| anyhow!("No body found"))
34}
35
36/// Same as `fetch_http`, but with additional headers and query support.
37pub fn fetch_http_with_options(host: &str, path: &str, options: RequestOptions) -> Result<String> {
38    let mut stream = TcpStream::connect((host, 80))?;
39
40    let url = build_url("", path, &options.query);
41    let headers = format_headers(host, &options);
42    let request = format!("GET {} HTTP/1.1\r\n{}\r\n\r\n", url, headers);
43
44    stream.write_all(request.as_bytes())?;
45
46    let mut response = String::new();
47    stream.read_to_string(&mut response)?;
48
49    response
50        .split("\r\n\r\n")
51        .nth(1)
52        .map(|body| body.to_string())
53        .ok_or_else(|| anyhow!("No body found"))
54}
55
56/// Extracts the values of a given attribute from all occurrences of a tag.
57pub fn extract_attribute(text: &str, tag: &str, attr: &str) -> Result<Vec<String>> {
58    let mut results = Vec::new();
59    let open_tag = format!("<{tag}");
60    let attr_eq = format!("{attr}=\"");
61
62    let mut start = 0;
63    while let Some(tag_start) = text[start..].find(&open_tag) {
64        let pos = start + tag_start;
65        if let Some(attr_start) = text[pos..].find(&attr_eq) {
66            let val_start = pos + attr_start + attr_eq.len();
67            if let Some(val_end) = text[val_start..].find('"') {
68                results.push(text[val_start..val_start + val_end].to_string());
69                start = val_start + val_end;
70            } else {
71                break;
72            }
73        } else {
74            start = pos + open_tag.len();
75        }
76    }
77
78    Ok(results)
79}
80
81/// Extract all href links from <a> tags.
82pub fn extract_links(text: &str) -> Result<Vec<String>> {
83    extract_attribute(text, "a", "href")
84}
85
86/// Extract inner content of all tags named `tag`.
87pub fn extract_tag(text: &str, tag: &str) -> Result<Vec<String>> {
88    let mut results = Vec::new();
89    let tag_lower = tag.to_lowercase();
90
91    let mut start = 0;
92    while let Some(open_pos) = text[start..].find('<') {
93        let open_tag_start = start + open_pos;
94        if let Some(open_tag_end) = text[open_tag_start..].find('>') {
95            let tag_content_start = open_tag_start + open_tag_end + 1;
96
97            let open_tag_full = &text[open_tag_start + 1..open_tag_start + open_tag_end];
98            let tag_name = open_tag_full
99                .split_whitespace()
100                .next()
101                .unwrap_or("")
102                .to_lowercase();
103
104            if tag_name == tag_lower {
105                let close_tag = format!("</{tag}>", tag = tag);
106                if let Some(close_pos) = text[tag_content_start..].to_lowercase().find(&close_tag) {
107                    let content = &text[tag_content_start..tag_content_start + close_pos];
108                    results.push(content.trim().to_string());
109                    start = tag_content_start + close_pos + close_tag.len();
110                    continue;
111                }
112            }
113
114            start = open_tag_start + 1;
115        } else {
116            break;
117        }
118    }
119
120    Ok(results)
121}
122
123/// Fetch a page over HTTPS with default headers.
124pub fn fetch_https(host: &str, path: &str) -> Result<String> {
125    let tcp = TcpStream::connect((host, 443)).map_err(|e| anyhow!("TLS connect failed: {e}"))?;
126
127    let connector = TlsConnector::new()?;
128    let mut stream = connector.connect(host, tcp)?;
129
130    let request = format!("GET {path} HTTP/1.1\r\nHost: {host}\r\nConnection: close\r\n\r\n");
131    stream.write_all(request.as_bytes())?;
132
133    let mut response = String::new();
134    stream.read_to_string(&mut response)?;
135
136    response
137        .split("\r\n\r\n")
138        .nth(1)
139        .map(|body| body.to_string())
140        .ok_or_else(|| anyhow!("No body found"))
141}
142
143/// Same as `fetch_https`, but with custom headers and query support.
144pub fn fetch_https_with_options(host: &str, path: &str, options: RequestOptions) -> Result<String> {
145    let tcp = TcpStream::connect((host, 443))?;
146    let connector = TlsConnector::new()?;
147    let mut stream = connector.connect(host, tcp)?;
148
149    let url = build_url("", path, &options.query);
150    let headers = format_headers(host, &options);
151    let request = format!("GET {} HTTP/1.1\r\n{}\r\n\r\n", url, headers);
152
153    stream.write_all(request.as_bytes())?;
154
155    let mut response = String::new();
156    stream.read_to_string(&mut response)?;
157
158    response
159        .split("\r\n\r\n")
160        .nth(1)
161        .map(|body| body.to_string())
162        .ok_or_else(|| anyhow!("No body found"))
163}
164
165// pub fn fetch_http_(host: &str, path: &str) -> PyResult<String> {
166//     let mut stream = TcpStream::connect((host, 80))
167//         .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(e.to_string()))?;
168
169//     stream.set_read_timeout(Some(std::time::Duration::from_secs(5)))
170//         .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(e.to_string()))?;
171//     stream.set_write_timeout(Some(std::time::Duration::from_secs(5)))
172//         .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(e.to_string()))?;
173
174//     let request = format!("GET {path} HTTP/1.1\r\nHost: {host}\r\nConnection: close\r\n\r\n");
175
176//     stream.write_all(request.as_bytes())
177//         .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(e.to_string()))?;
178
179//     let mut response = String::new();
180//     stream.read_to_string(&mut response)
181//         .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(e.to_string()))?;
182
183//     // Check HTTP status
184//     let status_line = response.lines().next().unwrap_or_default();
185//     if !status_line.contains("200 OK") {
186//         return Err(PyErr::new::<pyo3::exceptions::PyException, _>(format!("HTTP error: {}", status_line)));
187//     }
188
189//     if let Some(body) = response.split("\r\n\r\n").nth(1) {
190//         Ok(body.to_string())
191//     } else {
192//         Err(PyErr::new::<pyo3::exceptions::PyException, _>("No body found"))
193//     }
194// }