rubot/src/bot/modules/weather.rs

use std::collections::BTreeMap;
use std::fs;
use std::io::{self, Read, Write};
use std::iter::Peekable;
use std::net::{TcpStream, ToSocketAddrs};
use std::path::{Path, PathBuf};
use std::str::Chars;
use std::time::Duration;

use native_tls::TlsConnector;

use crate::bot::module::{Action, Command, CommandSpec, Module};

const COMMANDS: &[CommandSpec] = &[
    CommandSpec { name: "add", usage: "add <location>", about: "save your weather location" },
    CommandSpec { name: "w", usage: "w [location]", about: "weather for your saved location" },
];

pub struct Weather {
    store: Store,
    path: PathBuf,
}

impl Weather {
    pub fn new(network: &str) -> Weather {
        let dir = std::env::var("RUSTBOT_DATA_DIR").unwrap_or_else(|_| ".".to_string());
        let safe: String = network
            .chars()
            .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
            .collect();
        Weather::with_path(PathBuf::from(dir).join(format!("weather.{safe}.json")))
    }

    pub fn with_path(path: PathBuf) -> Weather {
        Weather { store: Store::load(&path), path }
    }

    pub fn location_of(&self, nick: &str) -> Option<&str> {
        self.store.get(nick).map(String::as_str)
    }
}

impl Module for Weather {
    fn name(&self) -> &'static str {
        "weather"
    }

    fn commands(&self) -> &'static [CommandSpec] {
        COMMANDS
    }

    fn on_command(&mut self, cmd: &Command) -> Vec<Action> {
        match cmd.name {
            "add" => {
                if cmd.args.is_empty() {
                    return vec![Action::reply(format!("usage: {}add <location>", cmd.prefix))];
                }
                let loc = cmd.args.join(" ");
                self.store.set(cmd.sender, &loc);
                if let Err(e) = self.store.save(&self.path) {
                    crate::log(&format!("weather: could not save {}: {e}", self.path.display()));
                    return vec![Action::reply(format!("{}: could not save location", cmd.sender))];
                }
                vec![Action::reply(format!("{}: saved {loc} — try {}w", cmd.sender, cmd.prefix))]
            }
            "w" => {
                let loc = if !cmd.args.is_empty() {
                    cmd.args.join(" ")
                } else if let Some(l) = self.store.get(cmd.sender) {
                    l.clone()
                } else {
                    return vec![Action::reply(format!(
                        "{}: no location set — {}add <location> first",
                        cmd.sender, cmd.prefix
                    ))];
                };
                match fetch_weather(&loc) {
                    Ok(text) => vec![Action::reply(text)],
                    Err(e) => vec![Action::reply(format!("weather unavailable: {e}"))],
                }
            }
            _ => vec![],
        }
    }
}

struct Store {
    map: BTreeMap<String, String>,
}

impl Store {
    fn load(path: &Path) -> Store {
        match fs::read_to_string(path) {
            Ok(s) => Store { map: parse_db(&s) },
            Err(_) => Store { map: BTreeMap::new() },
        }
    }

    fn save(&self, path: &Path) -> io::Result<()> {
        fs::write(path, to_json(&self.map) + "\n")
    }

    fn get(&self, nick: &str) -> Option<&String> {
        self.map.get(&nick.to_lowercase())
    }

    fn set(&mut self, nick: &str, loc: &str) {
        self.map.insert(nick.to_lowercase(), loc.to_string());
    }
}

fn fetch_weather(location: &str) -> Result<String, String> {
    let path = format!("/{}?format=%l:+%C+%t+feels+%f+%w+%h&m", url_encode(location));
    let body = https_get("wttr.in", &path)?;
    let line: String = body
        .lines()
        .map(str::trim)
        .find(|l| !l.is_empty())
        .unwrap_or("")
        .chars()
        .take(250)
        .collect();
    if line.is_empty() {
        return Err("no data".to_string());
    }
    Ok(line)
}

fn https_get(host: &str, path: &str) -> Result<String, String> {
    let addr = (host, 443)
        .to_socket_addrs()
        .map_err(|e| e.to_string())?
        .next()
        .ok_or_else(|| "dns lookup failed".to_string())?;
    let stream = TcpStream::connect_timeout(&addr, Duration::from_secs(5)).map_err(|e| e.to_string())?;
    stream.set_read_timeout(Some(Duration::from_secs(5))).ok();
    stream.set_write_timeout(Some(Duration::from_secs(5))).ok();

    let connector = TlsConnector::new().map_err(|e| e.to_string())?;
    let mut tls = connector.connect(host, stream).map_err(|e| e.to_string())?;

    let req = format!(
        "GET {path} HTTP/1.0\r\nHost: {host}\r\nUser-Agent: curl/rubot\r\nAccept: text/plain\r\nConnection: close\r\n\r\n"
    );
    tls.write_all(req.as_bytes()).map_err(|e| e.to_string())?;

    let mut buf = Vec::new();
    tls.read_to_end(&mut buf).map_err(|e| e.to_string())?;
    let text = String::from_utf8_lossy(&buf);
    match text.split_once("\r\n\r\n") {
        Some((_, body)) => Ok(body.to_string()),
        None => Err("malformed response".to_string()),
    }
}

fn url_encode(s: &str) -> String {
    let mut out = String::new();
    for b in s.bytes() {
        match b {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => out.push(b as char),
            _ => out.push_str(&format!("%{b:02X}")),
        }
    }
    out
}

pub fn to_json(map: &BTreeMap<String, String>) -> String {
    let items: Vec<String> = map
        .iter()
        .map(|(k, v)| format!("\"{}\":\"{}\"", json_escape(k), json_escape(v)))
        .collect();
    format!("{{{}}}", items.join(","))
}

pub fn parse_db(s: &str) -> BTreeMap<String, String> {
    let mut map = BTreeMap::new();
    let mut chars = s.chars().peekable();
    skip_ws(&mut chars);
    if chars.next() != Some('{') {
        return map;
    }
    loop {
        skip_ws(&mut chars);
        match chars.peek() {
            Some('"') => {}
            Some(',') => {
                chars.next();
                continue;
            }
            _ => break,
        }
        let Some(key) = parse_string(&mut chars) else { break };
        skip_ws(&mut chars);
        if chars.next() != Some(':') {
            break;
        }
        skip_ws(&mut chars);
        let Some(value) = parse_string(&mut chars) else { break };
        map.insert(key, value);
    }
    map
}

fn skip_ws(chars: &mut Peekable<Chars>) {
    while let Some(c) = chars.peek() {
        if c.is_whitespace() {
            chars.next();
        } else {
            break;
        }
    }
}

fn parse_string(chars: &mut Peekable<Chars>) -> Option<String> {
    if chars.next() != Some('"') {
        return None;
    }
    let mut out = String::new();
    while let Some(c) = chars.next() {
        match c {
            '"' => return Some(out),
            '\\' => match chars.next()? {
                '"' => out.push('"'),
                '\\' => out.push('\\'),
                '/' => out.push('/'),
                'n' => out.push('\n'),
                'r' => out.push('\r'),
                't' => out.push('\t'),
                'b' => out.push('\u{08}'),
                'f' => out.push('\u{0C}'),
                'u' => {
                    let mut code = 0u32;
                    for _ in 0..4 {
                        code = code * 16 + chars.next()?.to_digit(16)?;
                    }
                    if let Some(ch) = char::from_u32(code) {
                        out.push(ch);
                    }
                }
                other => out.push(other),
            },
            c => out.push(c),
        }
    }
    None
}

fn json_escape(s: &str) -> String {
    let mut out = String::with_capacity(s.len() + 2);
    for c in s.chars() {
        match c {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)),
            c => out.push(c),
        }
    }
    out
}