encryption_macros_encryption/
lib.rs

1use encryption_macros_utils::xor;
2use litrs::StringLit;
3use proc_macro::{Delimiter, Group, TokenStream, TokenTree};
4use quote::quote;
5
6/// xor a single string literal declaratively
7///
8/// ## Example
9/// ```
10/// use encryption_macros::e;
11///
12/// fn main() {
13///     let secret_string = e!{"piss you can not read this from the raw binary"};
14///     println!("our secret string: \"{}\"", secret_string)
15/// }
16/// ```
17
18#[proc_macro]
19pub fn e(token_stream: TokenStream) -> TokenStream {
20    let string_lit = StringLit::try_from(
21        token_stream
22            .into_iter()
23            .collect::<Vec<_>>()
24            .first()
25            .expect("there is no first token in the token stream"),
26    )
27    .expect("the token is not a string literal");
28    generate_decode_scope(encoded_literal(string_lit.value())).into()
29}
30
31/// xor all string literals in the provided scope
32///
33/// ## Example
34/// ```
35/// use encryption_macros::encrypt_strings;
36///
37/// encrypt_strings!{
38///     fn main() {
39///         println!("everything in this scope gets encrypted, {}", "even this")
40///     }
41/// }
42/// ```
43
44static FORMAT_ARGS_IDENTIFIERS: &[&'static str] = &["print", "println", "format"];
45
46#[proc_macro]
47pub fn encrypt_strings(token_stream: TokenStream) -> TokenStream {
48    parse_scope(token_stream, false)
49}
50
51/// xor all strings in a decorated function
52///
53/// ## Example
54/// ```
55/// use encryption_macros::encrypt_all_strings;
56///
57/// #[encrypt_all_strings]
58/// fn main() {
59///     println!("everything in this function gets encrypted, {}", "even this")
60/// }
61/// ```
62
63#[proc_macro_attribute]
64pub fn encrypt_all_strings(_metadata: TokenStream, token_stream: TokenStream) -> TokenStream {
65    encrypt_strings(token_stream)
66}
67
68fn parse_scope(token_stream: TokenStream, mut format_arg_literal: bool) -> TokenStream {
69    let mut new_stream = TokenStream::new();
70    let mut token_stream = token_stream.into_iter().peekable();
71    while let Some(tokentree) = token_stream.next() {
72        match tokentree {
73            TokenTree::Group(group) => new_stream.extend([TokenTree::from(Group::new(
74                group.delimiter(),
75                parse_scope(group.stream(), format_arg_literal),
76            ))]),
77            TokenTree::Literal(literal) => match StringLit::try_from(&literal) {
78                Ok(literal) if literal.value().len() > 0 => {
79                    if format_arg_literal {
80                        new_stream.extend(parse_format_args_literal(literal.value()));
81                        format_arg_literal = false;
82                    } else {
83                        new_stream.extend::<TokenStream>(
84                            generate_decode_scope(encoded_literal(literal.value())).into(),
85                        );
86                    }
87                }
88                _ => new_stream.extend([TokenTree::Literal(literal)]),
89            },
90            TokenTree::Ident(ident)
91                if FORMAT_ARGS_IDENTIFIERS.contains(&ident.to_string().as_str()) =>
92            {
93                if let Some(TokenTree::Punct(punct)) = token_stream.peek() {
94                    if punct.as_char() == '!' {
95                        format_arg_literal = true;
96                    }
97                }
98                new_stream.extend([TokenTree::Ident(ident)])
99            }
100            TokenTree::Punct(punct) if punct.as_char() == '#' => {
101                new_stream.extend([TokenTree::Punct(punct)]);
102                match token_stream.next() {
103                    Some(TokenTree::Group(group)) if group.delimiter() == Delimiter::Bracket => {
104                        new_stream.extend([TokenTree::Group(group)])
105                    }
106                    Some(token_tree) => new_stream.extend([token_tree]),
107                    None => break,
108                }
109            }
110            t => new_stream.extend([t]),
111        }
112    }
113    new_stream
114}
115
116fn parse_format_args_literal(literal: &str) -> TokenStream {
117    let mut start = 0;
118    let mut format_args = String::new();
119    let mut string_literals = Vec::new();
120    
121    loop {
122        if let Some(new_start) = literal[start..].find('{') {
123            if let Some(end) = literal[start + new_start..].find('}') {
124                let real_start = start + new_start;
125                string_literals.push(encoded_literal(&literal[..real_start]));
126                format_args.push_str("{}");
127                format_args.push_str(&literal[real_start..real_start + end + 1]);
128                start += new_start + end + 1;
129            } else {
130                break;
131            }
132        } else {
133            break;
134        }
135    }
136
137    if start != literal.len() - 1 {
138        format_args.push_str("{}");
139        string_literals.push(encoded_literal(&literal[start..]));
140    }
141
142    let decryption_scopes = string_literals
143        .iter()
144        .map(|string_literal| generate_decode_scope(string_literal.to_string()));
145
146    quote! {
147        #format_args,
148        #(#decryption_scopes),*
149    }
150    .into()
151}
152
153fn encoded_literal(l: &str) -> String {
154    let mut bytes = l.as_bytes().to_owned();
155    xor(&mut bytes);
156    hex::encode(bytes)
157}
158
159fn generate_decode_scope(hex_encoded_bytes: String) -> quote::__private::TokenStream {
160    quote! {
161        {
162            let mut bytes = encryption_macros::decode(#hex_encoded_bytes).unwrap();
163            encryption_macros::xor(&mut bytes);
164            String::from_utf8(bytes).unwrap()
165        }
166    }
167}
encryption_macros_encryption/lib.rs

encryption_macros_encryption/
lib.rs
