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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
#![forbid(unsafe_code)]
#![warn(
    clippy::all,
    clippy::pedantic,
    clippy::nursery,
    rustdoc::broken_intra_doc_links,
    missing_docs
)]

//! # `muddy_macro`
//!
//! This crate provides macros that are reexported by the [muddy](https://github.com/orph3usLyre/muddy-waters/tree/main/muddy) crate and should **not** be used
//! outside of the context of that crate.
//!
//!

use chacha20poly1305::{
    aead::{KeyInit, OsRng},
    ChaCha20Poly1305, Key,
};
use once_cell::sync::Lazy;
use proc_macro::TokenStream;
use proc_macro2::{Ident, Span};
use rand::{
    distributions::{Alphanumeric, DistString},
    Rng,
};
use syn::{parse_macro_input, LitStr};

mod internal;
mod meta;

use internal::{build_obfuscation_mod, build_static_obfuscation, encrypt_string_tokens};
use meta::{KeyMode, NonObfuscatedText, NonObfuscatedTexts};

/// Used to generate the key at buildtime
///
/// Kept seperately to embed in the target binary
pub(crate) static KEY: Lazy<Key> = Lazy::new(|| ChaCha20Poly1305::generate_key(&mut OsRng));

/// Used to generate text encryptions at build time
pub(crate) static ENCRYPTION: Lazy<ChaCha20Poly1305> = Lazy::new(|| ChaCha20Poly1305::new(&KEY));

/// Random string identifiers for the key and cipher
pub(crate) static IDENTS: Lazy<(String, String)> = Lazy::new(|| {
    static A: u8 = b'A';
    static Z: u8 = b'Z';
    (
        format!(
            "{}{}",
            OsRng.gen_range(A..=Z) as char,
            Alphanumeric.sample_string(&mut OsRng, 16).to_uppercase()
        ),
        format!(
            "{}{}",
            OsRng.gen_range(A..=Z) as char,
            Alphanumeric.sample_string(&mut OsRng, 16).to_uppercase()
        ),
    )
});

static ERROR_MESSAGE: &str = "Encountered encryption error";
static DEFAULT_ENV: &str = "MUDDY";

pub(crate) type Result<T> = std::result::Result<T, chacha20poly1305::Error>;

// NOTE: All doc tests in this crate are marked with `ignore` since the `muddy_macros` cannot work
// without the `muddy` crate as a wrapper

#[proc_macro]
/// Initialization macro that must be called at the crate root.
/// This sets up the scaffolding for the lazy decryption at runtime.
///
/// # Modes
///
/// Optional values that can be provided are:
/// - "embed" (default)
/// - "env"
///
/// ### "embed"
/// If the "embed" mode is chosen, the key will be embedded in the binary with minor obfuscation
/// (`XORed` with a random array).
///
/// ### "env"
/// If "env" is provided, the key will not be embedded in the binary.
/// An additional value may be provided to set the env key identifier:
///
/// ```ignore
///
/// muddy_init!("env", "MY_KEY");
/// ```
///
/// Or at buildtime:  
///
/// ```ignore
///
/// muddy_init!("env");
///
/// // run with `MUDDY='MY_KEY_NAME' cargo b`
///
/// ```
///
///
pub fn muddy_init(input: TokenStream) -> TokenStream {
    let keymode: KeyMode = parse_macro_input!(input as KeyMode);
    let key_ident = Ident::new(&IDENTS.0, Span::call_site());
    let cipher_ident = Ident::new(&IDENTS.1, Span::call_site());
    build_obfuscation_mod(&key_ident, &cipher_ident, keymode)
}

#[proc_macro]
/// Obfuscates a literal text.
/// (aka [`m!`])
///
/// # Example
///
/// ```ignore
/// println!("{}", muddy_str!("my text"));
/// ```
///
pub fn muddy_str(input: TokenStream) -> TokenStream {
    let text = parse_macro_input!(input as LitStr);
    let Ok(out) = encrypt_string_tokens(&text.value()) else {
        syn::Error::new(text.span(), ERROR_MESSAGE).to_compile_error();
        return TokenStream::new();
    };
    out
}

#[proc_macro]
/// Obfuscates a literal text.
/// (aka [`muddy_str!`])
///
/// # Example
///
/// ```ignore
/// println!("{}", m!("my text"));
/// ```
///
pub fn m(input: TokenStream) -> TokenStream {
    muddy_str(input)
}

#[proc_macro_attribute]
/// Obfuscates a static string expression as an attribute macro.
///
/// # Example
///
/// ```ignore
/// #[muddy]
/// static MY_STR: &str = "my text";
/// ```
///
pub fn muddy(_args: TokenStream, input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as NonObfuscatedText);
    let Ok(out) = build_static_obfuscation(&input) else {
        syn::Error::new(input.text.span(), ERROR_MESSAGE).to_compile_error();
        return TokenStream::new();
    };
    out
}

#[proc_macro]
/// Obfuscates multiple static string expressions.
///
/// # Example
///
/// ```ignore
/// muddy_all! {
///     pub static MY_FIRST_STR: &str = "my text";
///     static MY_SECOND_STR: &str = "my second text";
/// }
/// ```
///
pub fn muddy_all(input: TokenStream) -> TokenStream {
    let non_obfuscated = parse_macro_input!(input as NonObfuscatedTexts);
    let mut output: TokenStream = TokenStream::new();
    let Ok(iter): Result<Vec<TokenStream>> = non_obfuscated
        .texts
        .iter()
        .map(build_static_obfuscation)
        .collect()
    else {
        syn::Error::new(Span::call_site(), ERROR_MESSAGE).to_compile_error();
        return TokenStream::new();
    };
    output.extend(iter);
    output
}