1use std::{
2 borrow::Cow,
3 collections::{HashMap, HashSet},
4 fmt, io,
5};
6
7use colored::Colorize;
8
9use crate::{
10 db::UserData,
11 scanner::{EtcHostScanner, HostScanner},
12 utils::{get_host_from_url_or_host, is_valid_url, sha256},
13};
14
15pub struct App<'a, O: io::Write, E: io::Write> {
16 pub(crate) block: HashSet<Cow<'a, str>>,
17 pub(crate) data: UserData<'a>,
18 pub(crate) etc_hosts: Cow<'a, str>,
19 pub(crate) stdout: &'a mut O,
20 pub(crate) stderr: &'a mut E,
21}
22
23impl<'a, O: io::Write, E: io::Write> App<'a, O, E> {
24 pub fn new<R: io::Read>(
25 etc_hosts: &'a str,
26 db: Option<R>,
27 stdout: &'a mut O,
28 stderr: &'a mut E,
29 ) -> io::Result<Self> {
30 let user_db = db
31 .map(|v| UserData::from_read(v).unwrap_or_default())
32 .unwrap_or_default();
33 let mut block = HashSet::with_capacity(etc_hosts.len() / 20);
34 for i in EtcHostScanner::from(etc_hosts) {
35 block.insert(Cow::Borrowed(i));
36 }
37 Ok(Self {
38 block,
39 data: user_db,
40 etc_hosts: Cow::Borrowed(etc_hosts),
41 stdout,
42 stderr,
43 })
44 }
45 fn eprintln_invalid_host_or_url<T: fmt::Display>(&mut self, value: T) {
46 let _ = writeln!(
47 self.stderr,
48 "ERROR: invalid host or url: {}",
49 value.to_string().italic().bold().red(),
50 );
51 let _ = self.stderr.flush();
52 }
53 fn eprintln_url<T: fmt::Display>(&mut self, value: T) {
54 let _ = writeln!(
55 self.stderr,
56 "ERROR: invalid url: {}",
57 value.to_string().italic().bold().red()
58 );
59 }
60 #[inline]
61 pub fn get_sources(&self) -> &HashMap<Cow<'a, str>, [u8; 32]> {
62 &self.data.sources
63 }
64 pub fn add_allow<T: Iterator<Item = &'a str>>(&mut self, args: T) {
65 for i in args {
66 if let Some(v) = get_host_from_url_or_host(i) {
67 let val = Cow::Borrowed(v);
68 self.block.remove(&val);
69 self.data.insert_allow(val);
70 } else {
71 self.eprintln_invalid_host_or_url(i);
72 }
73 }
74 }
75 pub fn rm_allow<T: Iterator<Item = &'a str>>(&mut self, args: T) {
76 for i in args {
77 if let Some(host) = get_host_from_url_or_host(i) {
78 self.data.remove_allow(host);
79 } else {
80 self.eprintln_invalid_host_or_url(i);
81 }
82 }
83 }
84 pub fn add_block<T: Iterator<Item = &'a str>>(&mut self, args: T) {
85 for i in args {
86 if let Some(v) = get_host_from_url_or_host(i) {
87 let val = Cow::Borrowed(v);
88 self.data.insert_block(val);
89 } else {
90 self.eprintln_invalid_host_or_url(i);
91 }
92 }
93 }
94 pub fn rm_block<T: Iterator<Item = &'a str>>(&mut self, args: T) {
95 for i in args {
96 if let Some(v) = get_host_from_url_or_host(i) {
97 self.data.remove_block(v);
98 } else {
99 self.eprintln_invalid_host_or_url(i);
100 }
101 }
102 }
103 pub fn add_redirect<T: Iterator<Item = (&'a str, &'a str)>>(&mut self, args: T) {
104 for i in args {
105 if let Some(to) = get_host_from_url_or_host(i.0) {
106 if let Some(from) = get_host_from_url_or_host(i.1) {
107 let to = Cow::Borrowed(to);
108 let from = Cow::Borrowed(from);
109 self.block.remove(&to);
110 self.block.remove(&from);
111 self.data.insert_redirect((to, from));
112 } else {
113 self.eprintln_invalid_host_or_url(i.1);
114 }
115 } else {
116 self.eprintln_invalid_host_or_url(i.0);
117 }
118 }
119 }
120 pub fn rm_redirect<T: Iterator<Item = &'a str>>(&mut self, args: T) {
121 for i in args {
122 if let Some(v) = get_host_from_url_or_host(i) {
123 self.data.redirect.remove(&Cow::Borrowed(v));
124 } else {
125 self.eprintln_invalid_host_or_url(i);
126 }
127 }
128 }
129 pub fn add_sources<T: Iterator<Item = &'a str>>(&mut self, args: T) {
130 for i in args {
131 if is_valid_url(i) {
132 self.data.insert_sources(Cow::Borrowed(i));
133 } else {
134 self.eprintln_url(i);
135 }
136 }
137 }
138 pub fn rm_sources<T: Iterator<Item = &'a str>>(&mut self, args: T) {
139 for i in args {
140 if is_valid_url(i) {
141 self.data.remove_sources(i);
142 }
143 }
144 }
145 #[allow(clippy::result_large_err)]
146 fn download<T: AsRef<str>>(url: T) -> Result<String, ureq::Error> {
147 ureq::get(url.as_ref()).call()?.body_mut().read_to_string()
148 }
149 pub fn get_update(&mut self) -> Vec<(String, String, [u8; 32])> {
150 let mut v = Vec::with_capacity(self.data.sources.len());
151 for (url, hash) in self.data.sources.iter() {
152 let _ = writeln!(self.stdout, "Checking: {}", url.yellow());
153 match Self::download(url) {
154 Ok(s) => {
155 let new_hash = sha256(&s);
156 if &new_hash != hash {
157 let _ = writeln!(self.stdout, "...Update Available...\n");
158 v.push((s, url.to_string(), new_hash));
159 continue;
160 } else {
161 let _ = writeln!(self.stdout, "...Update Not Available...\n");
162 v.push((s, url.to_string(), *hash));
163 }
164 }
165 Err(e) => {
166 let _ = writeln!(self.stderr, "{e}");
167 }
168 }
169 }
170 v
171 }
172 pub fn print_etc_hosts(&mut self) -> io::Result<()> {
173 for line in self.etc_hosts.lines() {
174 self.stdout.write_all(line.as_bytes())?;
175 self.stdout.write_all(b"\n")?;
176 }
177 self.stdout.flush()
178 }
179 pub fn get_update_fource(&mut self) -> Vec<(String, String, [u8; 32])> {
180 let mut v = Vec::with_capacity(self.data.sources.len());
181 for (url, _) in self.data.sources.iter() {
182 let _ = writeln!(self.stdout, "Downloading From: {}", url.yellow());
183 match Self::download(url) {
184 Ok(s) => {
185 let hash = sha256(&s);
186 v.push((s, url.to_string(), hash));
187 }
188 Err(e) => {
189 let _ = writeln!(self.stderr, "{e}");
190 }
191 }
192 }
193 v
194 }
195 pub fn apply_update(&mut self, update: &'a [(String, String, [u8; 32])]) {
196 let mut est_len = 0;
197 self.block.clear();
198 for (data, _, _) in update.iter() {
199 est_len += data.len();
200 }
201 est_len /= 20;
202 if self.block.capacity() < est_len {
203 let _ = self.block.try_reserve(est_len - self.block.capacity());
204 }
205 let mut update_flag = false;
206 for (data, url, hash) in update.iter() {
207 for host in HostScanner::from(data.as_str()) {
208 self.block.insert(Cow::Borrowed(host));
209 }
210 if let Some(h) = self.data.sources.get_mut(&Cow::Borrowed(url.as_str())) {
211 if h != hash {
212 update_flag = true;
213 }
214 *h = *hash;
215 }
216 }
217 if update_flag {
218 let _ = self.stdout.write_all(b".....Update Success.....\n");
219 }
220 }
221
222 pub fn clear_host(&mut self) {
223 self.block.clear();
224 }
225 pub fn export<W: io::Write>(&mut self, w: &mut W) -> io::Result<()> {
226 self.data.write(w)
227 }
228 pub fn save_1<W1: io::Write, W2: io::Write>(
231 &mut self,
232 w1: &mut W1,
233 w2: &mut W2,
234 ) -> io::Result<()> {
235 self.save(|| (w1, w2))
236 }
237 pub fn save<W1: io::Write, W2: io::Write, F: FnOnce() -> (W1, W2)>(
240 &mut self,
241 f: F,
242 ) -> io::Result<()> {
243 let mut block = HashSet::with_capacity(self.block.len() + self.data.block.len());
244 for i in self.block.iter() {
245 block.insert(i.as_bytes());
246 }
247 for i in self.data.block.iter() {
248 block.insert(i.as_bytes());
249 }
250 for i in self.data.allow.iter() {
251 block.remove(i.as_bytes());
252 }
253 let mut v = Vec::with_capacity(block.len());
254 v.extend(block);
255 v.sort();
256 let mut redirect = Vec::with_capacity(self.data.redirect.len());
257 for i in self.data.redirect.iter() {
258 redirect.push((i.1.as_bytes(), i.0.as_bytes()));
259 }
260 redirect.sort();
261 let (mut etc, mut data) = f();
262 write_etc_host(v, redirect, self.etc_hosts.as_ref(), &mut etc)?;
263 self.data.write(&mut data)?;
264 self.stdout.write_all(b".....Saved Changes.....\n")?;
265 Ok(())
266 }
267 pub fn restore_etc_hosts<W: io::Write>(etc_hosts: &str, w: &mut W) -> io::Result<()> {
268 let mut iter = etc_hosts.lines();
269 while let Some(line) = iter.next() {
270 match line {
271 "#host-rs-beg#" => {
272 for line in iter.by_ref() {
273 if line == "#host-rs-end#" {
274 break;
275 }
276 }
277 }
278 "#r-host-rs-beg#" => {
279 for line in iter.by_ref() {
280 if line == "#r-host-rs-end#" {
281 break;
282 }
283 }
284 }
285 v => {
286 w.write_all(v.as_bytes())?;
287 w.write_all(b"\n")?;
288 }
289 }
290 }
291 w.flush()
292 }
293 pub fn clear_data(&mut self) {
294 self.data.clear();
295 }
296 pub fn restore_data<W: io::Write>(w: &mut W) -> io::Result<()> {
297 UserData::default().write(w)
298 }
299}
300
301fn write_etc_host<W>(
302 block: Vec<&[u8]>,
303 redirect: Vec<(&[u8], &[u8])>,
304 old_etc: &str,
305 etc_file: &mut W,
306) -> io::Result<()>
307where
308 W: io::Write,
309{
310 let block_start = b"#host-rs-beg#";
311 let block_end = b"#host-rs-end#";
312 let r_start = b"#r-host-rs-beg#";
313 let r_end = b"#r-host-rs-end#";
314 let mut old_etc = old_etc.lines().map(|v| v.as_bytes());
315 while let Some(line) = old_etc.next() {
316 match line {
317 v if v == block_start => {
318 for line in old_etc.by_ref() {
319 if line == block_end {
320 break;
321 }
322 }
323 }
324 v if v == r_start => {
325 for line in old_etc.by_ref() {
326 if line == r_end {
327 break;
328 }
329 }
330 }
331 v => {
332 etc_file.write_all(v)?;
333 etc_file.write_all(b"\n")?;
334 }
335 }
336 }
337 etc_file.write_all(block_start)?;
338 etc_file.write_all(b"\n")?;
339 for i in block {
340 etc_file.write_all(b"0.0.0.0 ")?;
341 etc_file.write_all(i.as_ref())?;
342 etc_file.write_all(b"\n")?;
343 }
344 etc_file.write_all(block_end)?;
345 etc_file.write_all(b"\n")?;
346 etc_file.write_all(r_start)?;
347 etc_file.write_all(b"\n")?;
348 for i in redirect {
349 etc_file.write_all(i.0.as_ref())?;
350 etc_file.write_all(b" ")?;
351 etc_file.write_all(i.1.as_ref())?;
352 etc_file.write_all(b"\n")?;
353 }
354 etc_file.write_all(r_end)?;
355 etc_file.write_all(b"\n")?;
356 etc_file.flush()?;
357 Ok(())
358}