i3status_rs/blocks/vpn/
mullvad.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
use regex::Regex;
use std::process::Stdio;
use tokio::process::Command;

use crate::blocks::prelude::*;
use crate::util::country_flag_from_iso_code;

use super::{Driver, Status};

pub struct MullvadDriver {
    regex_country_code: Regex,
}

impl MullvadDriver {
    pub async fn new() -> MullvadDriver {
        MullvadDriver {
            regex_country_code: Regex::new("Connected to ([a-z]{2}).*, ([A-Z][a-z]*).*\n").unwrap(),
        }
    }

    async fn run_network_command(arg: &str) -> Result<()> {
        let code = Command::new("mullvad")
            .args([arg])
            .stdin(Stdio::null())
            .stdout(Stdio::null())
            .spawn()
            .error(format!("Problem running mullvad command: {arg}"))?
            .wait()
            .await
            .error(format!("Problem running mullvad command: {arg}"))?;

        if code.success() {
            Ok(())
        } else {
            Err(Error::new(format!(
                "mullvad command failed with nonzero status: {code:?}"
            )))
        }
    }
}

#[async_trait]
impl Driver for MullvadDriver {
    async fn get_status(&self) -> Result<Status> {
        let stdout = Command::new("mullvad")
            .args(["status"])
            .output()
            .await
            .error("Problem running mullvad command")?
            .stdout;

        let status = String::from_utf8(stdout).error("mullvad produced non-UTF8 output")?;

        if status.contains("Disconnected") {
            return Ok(Status::Disconnected);
        } else if status.contains("Connected") {
            let (country_flag, country) = self
                .regex_country_code
                .captures_iter(&status)
                .next()
                .map(|capture| {
                    let country_code = capture[1].to_uppercase();
                    let country = capture[2].to_owned();
                    let country_flag = country_flag_from_iso_code(&country_code);
                    (country_flag, country)
                })
                .unwrap_or_default();

            return Ok(Status::Connected {
                country,
                country_flag,
            });
        }
        Ok(Status::Error)
    }

    async fn toggle_connection(&self, status: &Status) -> Result<()> {
        match status {
            Status::Connected { .. } => Self::run_network_command("disconnect").await?,
            Status::Disconnected => Self::run_network_command("connect").await?,
            Status::Error => (),
        }
        Ok(())
    }
}