Add obfuscated logger.

unshell-obfuscate/src/encrypt.rs

@@ -19,6 +19,11 @@ const STATIC_IV: [u8; 16] = [
     0x6d, 0x79, 0x5f, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x69, 0x76, 0x5f, 0x30, 0x31, 0x32,
 ];
 
+fn pkcs7_padded_length(input_len: usize) -> usize {
+    let block_size = 16;
+    ((input_len / block_size) + 1) * block_size
+}
+
 pub fn get_obfuscated_symbol_name(input: &str) -> String {
     // 1. Get the key from the environment
     // let key_str =
@@ -36,7 +41,9 @@ pub fn get_obfuscated_symbol_name(input: &str) -> String {
     let mut plaintext = input.to_string();
     let plaintext = unsafe { plaintext.as_bytes_mut() };
 
-    let mut buf = [0u8; 48];
+    let buf_len = pkcs7_padded_length(plaintext.len());
+    let mut buf: Vec<u8> = vec![0; buf_len];
+
     buf[..plaintext.len()].copy_from_slice(plaintext);
     let ciphertext = cipher
         .encrypt_padded_mut::<Pkcs7>(&mut buf, plaintext.len())

unshell-obfuscate/src/format_helper.rs

@@ -0,0 +1,97 @@
+use syn::parse::{Parse, ParseStream};
+use syn::{Expr, Lit, Token};
+
+pub struct PrintlnArgs {
+    pub format_str: String,
+    pub args: Vec<Expr>,
+}
+
+impl Parse for PrintlnArgs {
+    fn parse(input: ParseStream) -> syn::Result<Self> {
+        let format_expr: Expr = input.parse()?;
+
+        let format_str = match format_expr {
+            Expr::Lit(ref lit) => {
+                if let Lit::Str(ref s) = lit.lit {
+                    s.value()
+                } else {
+                    return Err(syn::Error::new_spanned(lit, "Expected string literal"));
+                }
+            }
+            _ => {
+                return Err(syn::Error::new_spanned(
+                    format_expr,
+                    "Expected string literal",
+                ));
+            }
+        };
+
+        let mut args = Vec::new();
+        while !input.is_empty() {
+            input.parse::<Token![,]>()?;
+            if input.is_empty() {
+                break;
+            }
+            args.push(input.parse()?);
+        }
+
+        Ok(PrintlnArgs { format_str, args })
+    }
+}
+
+#[derive(Debug)]
+pub enum FormatSegment {
+    Static(String),
+    Dynamic(String, usize), // format spec, arg index
+}
+
+pub fn parse_format_string(fmt: &str) -> Vec<FormatSegment> {
+    let mut segments = Vec::new();
+    let mut current_static = String::new();
+    let mut chars = fmt.chars().peekable();
+    let mut arg_idx = 0;
+
+    while let Some(ch) = chars.next() {
+        if ch == '{' {
+            if chars.peek() == Some(&'{') {
+                chars.next();
+                current_static.push('{');
+                continue;
+            }
+
+            // Save current static segment
+            if !current_static.is_empty() {
+                segments.push(FormatSegment::Static(current_static.clone()));
+                current_static.clear();
+            }
+
+            // Parse format spec
+            let mut spec = String::new();
+            while let Some(&next_ch) = chars.peek() {
+                if next_ch == '}' {
+                    chars.next();
+                    break;
+                }
+                spec.push(chars.next().unwrap());
+            }
+
+            segments.push(FormatSegment::Dynamic(spec, arg_idx));
+            arg_idx += 1;
+        } else if ch == '}' {
+            if chars.peek() == Some(&'}') {
+                chars.next();
+                current_static.push('}');
+            } else {
+                current_static.push(ch);
+            }
+        } else {
+            current_static.push(ch);
+        }
+    }
+
+    if !current_static.is_empty() {
+        segments.push(FormatSegment::Static(current_static));
+    }
+
+    segments
+}

unshell-obfuscate/src/lib.rs

@@ -1,14 +1,16 @@
 #![feature(proc_macro_quote)]
+#![feature(proc_macro_span)]
 
 use proc_macro::TokenStream;
-
 use quote::quote;
-
-use syn::{ItemFn, parse_macro_input};
+use syn::{Expr, ItemFn, LitStr, parse_macro_input};
 
 #[cfg(feature = "obfuscate")]
 mod encrypt;
 
+mod format_helper;
+use format_helper::*;
+
 // Put all encrypt-related dependencies in a module, so they are easier to use with the feature flag
 #[cfg(feature = "obfuscate")]
 mod obs_deps {
@@ -132,3 +134,107 @@ pub fn obs(input: TokenStream) -> TokenStream {
 
     TokenStream::from(obfuscated_expansion)
 }
+
+// #[proc_macro]
+// pub fn file_literal(_input: TokenStream) -> TokenStream {
+//     // let input = input.into_iter().collect::<Vec<_>>();
+//     // if input.len() != 1 {
+//     //     let msg = format!("expected exactly one input token, got {}", input.len());
+//     //     return quote! { compile_error!(#msg) }.into();
+//     // }
+
+//     let string = file!();
+//     let lit_str = LitStr::new(string, proc_macro2::Span::call_site());
+
+//     // let string_lit = match LitStr::try_from(&input) {
+//     //     // Error if the token is not a string literal
+//     //     Err(e) => return e.to_compile_error(),
+//     //     Ok(lit) => lit,
+//     // };
+
+//     (quote! {
+//         #lit_str
+//     })
+//     .into()
+// }
+
+#[proc_macro]
+pub fn file_symbol(_input: TokenStream) -> TokenStream {
+    // Get the call site span to extract file information
+    let span = proc_macro::Span::call_site();
+    let source_file = span.source();
+    let file_path = source_file.file();
+    let line_num = source_file.line();
+    let concatted = format!("{}:{}", file_path, line_num);
+
+    // Return as a string literal
+    let output = quote! {
+        unshell_obfuscate::symbol!(#concatted)
+    };
+    // let output = quote! {
+    //     #concatted
+    // };
+    output.into()
+}
+
+#[proc_macro]
+pub fn format_obs(input: TokenStream) -> TokenStream {
+    let PrintlnArgs { format_str, args } = parse_macro_input!(input as PrintlnArgs);
+
+    let segments = parse_format_string(&format_str);
+
+    if segments.is_empty() {
+        return quote! {
+            print!("\n")
+        }
+        .into();
+    }
+
+    let mut parts = Vec::new();
+
+    for segment in segments {
+        match segment {
+            FormatSegment::Static(text) => {
+                parts.push(quote! {
+                    unshell_obfuscate::symbol!(#text).to_string()
+                });
+            }
+            FormatSegment::Dynamic(spec, idx) => {
+                if idx >= args.len() {
+                    return syn::Error::new(
+                        proc_macro2::Span::call_site(),
+                        format!("argument {} is missing", idx),
+                    )
+                    .to_compile_error()
+                    .into();
+                }
+
+                let arg = &args[idx];
+                let fmt_spec = if spec.is_empty() {
+                    quote! { "{}" }
+                } else {
+                    let full_spec = format!("{{{}}}", spec);
+                    quote! { #full_spec }
+                };
+
+                // quote! {
+                //     println!(#fmt_spec, #arg);
+                // }
+                parts.push(quote! {
+                    format!(#fmt_spec, #arg)
+                });
+            }
+        }
+    }
+
+    (quote! {
+        {
+            let mut string = String::new();
+            #(
+                string.push_str(&#parts);
+            )*
+            string
+        }
+    })
+    .into()
+}