atsh-lib 0.1.2

autossh/@shell library is used to ssh login and allow automatically login without password
Documentation 
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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

use base64::{Engine as _, engine::general_purpose};
use chacha20poly1305::ChaCha20Poly1305;
use chacha20poly1305::aead::generic_array::GenericArray;
use chacha20poly1305::aead::generic_array::typenum::Unsigned;
use chacha20poly1305::aead::{Aead, AeadCore, KeyInit, OsRng};
use std::io::{Error, ErrorKind};
use std::sync::LazyLock;
use tracing::{debug, warn};

static ATSH_KEY: LazyLock<Option<String>> = LazyLock::new(|| {
    if let Ok(key) = std::env::var("ATSH_KEY") {
        debug!("`ATSH_KEY` found in environment variable");
        return Some(key);
    }
    if let Ok(key) = std::env::var("ASKEY") {
        warn!("💡 Deprecated `ASKEY` in next version and use `ATSH_KEY` instead");
        return Some(key);
    }
    // warn!("💥 export `ASKEY` to protect password! 💥");
    None
});

pub fn check_secure() -> Result<(), Error> {
    if ATSH_KEY.is_none() {
        Err(Error::new(
            ErrorKind::NotFound,
            "💥 Export `ATSH_KEY` to protect password",
        ))
    } else {
        Ok(())
    }
}

fn generate_key(key: Option<&str>) -> Vec<u8> {
    if key.is_none() {
        return ChaCha20Poly1305::generate_key(&mut OsRng).to_vec();
    }
    let key = key.unwrap();
    let key = key.as_bytes().to_vec();

    // the key must 32 bytes
    if key.len() > 32 {
        return key[..32].to_vec();
    } else if key.len() < 32 {
        // resize it
        let mut key = key;
        key.resize(32, 0);
        return key;
    }
    return key;
}

fn chacha_encrypt(cleartext: &str, key: &[u8]) -> Vec<u8> {
    let cipher = ChaCha20Poly1305::new(GenericArray::from_slice(key));
    let nonce = ChaCha20Poly1305::generate_nonce(&mut OsRng);
    let mut obsf = cipher
        .encrypt(&nonce, cleartext.as_bytes())
        .expect("encrypt failed by chacha20");
    obsf.splice(..0, nonce.iter().copied());
    obsf
}

fn chacha_decrypt(obsf: &[u8], key: &[u8]) -> String {
    type NonceSize = <ChaCha20Poly1305 as AeadCore>::NonceSize;
    let cipher = ChaCha20Poly1305::new(GenericArray::from_slice(key));
    let (nonce, ciphertext) = obsf.split_at(NonceSize::to_usize());
    let nonce = GenericArray::from_slice(nonce);
    let plaintext = cipher
        .decrypt(nonce, ciphertext)
        .expect("decrypt failed by chacha20");
    String::from_utf8(plaintext).unwrap()
}

pub fn encrypt(data: &str) -> String {
    if ATSH_KEY.is_none() {
        return data.to_string();
    }
    // log::debug!("we found `ASKEY` and will encrypt.");
    let key = generate_key(ATSH_KEY.as_deref());
    let obsf = chacha_encrypt(data, &key);
    general_purpose::STANDARD_NO_PAD.encode(obsf)
}

pub fn decrypt(data: impl AsRef<str>) -> String {
    let data = data.as_ref();
    if ATSH_KEY.is_none() {
        return data.to_string();
    }
    // log::debug!("we found `ASKEY` and will decrypt.");
    let obsf = general_purpose::STANDARD_NO_PAD
        .decode(data.as_bytes())
        .expect("decode failed by base64");
    let key = generate_key(ATSH_KEY.as_deref());
    chacha_decrypt(&obsf, &key)
}

// tests
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_chacha() {
        let key = generate_key(None);
        println!("chacha key: {:?}", key);
        let ciphertext = chacha_encrypt("plaintext message", &key);
        println!("encrypt: {:?}", ciphertext);
        let plaintext = chacha_decrypt(&ciphertext, &key);
        println!("decrypt: {:?}", plaintext);
        assert_eq!(plaintext, "plaintext message");
    }

    #[test]
    fn test_secure() {
        // set env
        let key = generate_key(None);
        println!("secure key: {:?}", key);
        let key = "32 bit key must, if not, we will resize it.";
        unsafe { std::env::set_var("ASKEY", key) };
        let key = generate_key(Some(key));
        println!("secure key: {:?}", key);
        let data = "hello world";
        let enc = encrypt(data);
        println!("encrypt: {:?}", enc);
        let dec = decrypt(&enc);
        println!("decrypt: {:?}", dec);
        assert_eq!(data, dec);
    }
}