psocks/src/rules.rs

mod item;

use anyhow::{Result, bail};
use item::Item;
use log::*;
use std::collections::{HashMap, HashSet};
use tokio::{fs, sync::RwLock};

pub struct Rules(RwLock<HashMap<usize, List>>);

impl Rules {
    pub async fn from_opt(list: &[String]) -> Result<Self> {
        let mut index = HashMap::with_capacity(list.len());
        assert!(
            index
                .insert(
                    0,
                    List {
                        alias: "session".into(),
                        items: HashSet::new(),
                        status: true
                    }
                )
                .is_none()
        );
        for (i, l) in list.iter().enumerate() {
            let mut items = HashSet::new();
            for line in if l.contains("://") {
                let response = reqwest::get(l).await?;
                let status = response.status();
                if status.is_success() {
                    response.text().await?
                } else {
                    bail!("Could not receive remote list `{l}`: `{status}`")
                }
            } else {
                fs::read_to_string(l).await?
            }
            .lines()
            {
                if line.starts_with("/") || line.starts_with("#") || line.is_empty() {
                    continue; // skip comments
                }
                if !items.insert(Item::from_line(line)) {
                    warn!("List `{l}` contains duplicated entry: `{line}`")
                }
            }
            assert!(
                index
                    .insert(
                        i + 1,
                        List {
                            alias: l.clone(),
                            items,
                            status: true // @TODO implement config file
                        }
                    )
                    .is_none()
            )
        }
        Ok(Self(RwLock::new(index)))
    }
    /// Check if rule is exist in the (allow) index
    pub async fn any(&self, value: &str) -> bool {
        self.0.read().await.values().any(|list| list.any(value))
    }
    /// Get total rules from the current session
    pub async fn total(&self, status: bool) -> u64 {
        self.0
            .read()
            .await
            .values()
            .filter(|list| list.status == status)
            .map(|list| list.total())
            .sum()
    }
    /// Allow given `rule`(in the session index)
    /// * return `false` if the `rule` is exist
    pub async fn allow(&self, rule: &str) -> bool {
        self.0
            .write()
            .await
            .get_mut(&0)
            .unwrap()
            .items
            .insert(Item::from_line(rule))
    }
    /// Block given `rule` (in the session index)
    pub async fn block(&self, rule: &str) {
        self.0
            .write()
            .await
            .get_mut(&0)
            .unwrap()
            .items
            .retain(|item| rule == item.as_str())
    }
    /// Return active rules
    pub async fn active(&self) -> Vec<String> {
        let mut rules: Vec<String> = self
            .0
            .read()
            .await
            .values()
            .filter(|list| list.status)
            .flat_map(|list| list.items.iter().map(|item| item.to_string()))
            .collect();
        rules.sort(); // HashSet does not keep the order
        rules
    }
    /// Return list references
    pub async fn lists(&self) -> Vec<ListEntry> {
        let this = self.0.read().await;
        let mut list = Vec::with_capacity(this.len());
        for l in this.iter() {
            list.push(ListEntry {
                id: *l.0,
                alias: l.1.alias.clone(),
            })
        }
        list
    }
    /// Return original list references
    pub async fn list(&self, id: &usize) -> Option<List> {
        self.0.read().await.get(id).cloned()
    }
    /// Change rule set status by list ID
    pub async fn enable(&self, list_id: &usize, status: bool) -> Option<()> {
        self.0
            .write()
            .await
            .get_mut(list_id)
            .map(|this| this.status = status)
    }
}

#[derive(serde::Serialize)]
pub struct ListEntry {
    pub id: usize,
    pub alias: String,
}

#[derive(serde::Serialize, Clone)]
pub struct List {
    pub alias: String,
    pub items: HashSet<Item>,
    pub status: bool,
}

impl List {
    /// Check if rule is exist in the items index
    pub fn any(&self, value: &str) -> bool {
        self.items.iter().any(|item| match item {
            Item::Exact(v) => v == value,
            Item::Ending(v) => value.ends_with(v),
        })
    }
    /// Get total rules in list
    pub fn total(&self) -> u64 {
        self.items.len() as u64
    }
}