Merge 203ec47 into 5414f78

cli/Cargo.toml

@@ -18,4 +18,5 @@ panopticon-avr = { path = "../avr" }
 panopticon-graph-algos = { path = "../graph-algos" }
 log = "0.3"
 env_logger = "0.3"
-colored = "1.5"
+termcolor = "0.3.2"
+atty = "0.2.2"

cli/src/display.rs

@@ -1,44 +1,63 @@
-use colored::*;
+use std::io::Write;
+use atty;
+use termcolor::{BufferWriter, ColorChoice, ColorSpec, WriteColor};
+use termcolor::Color::*;
 
-use panopticon_core::{Function, BasicBlock, Mnemonic, MnemonicFormatToken, Program, Rvalue};
+use panopticon_core::{Function, BasicBlock, Mnemonic, MnemonicFormatToken, Program, Rvalue, Result};
+
+macro_rules! color_bold {
+    ($fmt:ident, $color:ident, $str:expr) => ({
+    $fmt.set_color(ColorSpec::new().set_bold(true).set_fg(Some($color)))?;
+    write!($fmt, "{}", $str)?;
+    $fmt.reset()?;
+    })
+}
+
+macro_rules! color {
+    ($fmt:ident, $color:ident, $str:expr) => ({
+        $fmt.set_color(ColorSpec::new().set_fg(Some($color)))?;
+        write!($fmt, "{}", $str)?;
+        $fmt.reset()?;
+    })
+}
 
 /// Displays the function in a human readable format, using `program`
-pub fn display_function(function: &Function, program: &Program) -> String {
+pub fn print_function(function: &Function, program: &Program, always_color: bool) -> Result<()> {
+    let cc = if always_color || atty::is(atty::Stream::Stdout) { ColorChoice::Auto }else { ColorChoice::Never };
+    let writer = BufferWriter::stdout(cc);
+    let mut fmt = writer.buffer();
     let mut bbs = function.basic_blocks().collect::<Vec<&BasicBlock>>();
     bbs.sort_by(|bb1, bb2| bb1.area.start.cmp(&bb2.area.start));
-    let mut fmt = {
-        let &BasicBlock { ref area, .. } = function.entry_point();
-        format!("{:0>8x} <{}>:", area.start, function.name.bold().yellow())
-    };
+    write!(fmt, "{:0>8x} <", function.start())?;
+    color_bold!(fmt, Yellow, function.name);
+    writeln!(fmt, ">:")?;
     for bb in bbs {
-        fmt = format!("{}{}", fmt, display_basic_block(&bb, program));
+        display_basic_block(&mut fmt, &bb, program)?;
     }
-    fmt
+    writer.print(&fmt)?;
+    Ok(())
 }
 
 /// Displays the basic block in disassembly order, in human readable form, and looks up any functions calls in `program`
-pub fn display_basic_block(basic_block: &BasicBlock, program: &Program) -> String {
-    let seed = String::new();
-    let display = basic_block.mnemonics.iter().filter_map(|x| {
-        if x.opcode.starts_with("__") {
-            None
-        } else {
-            Some(x)
+pub fn display_basic_block<W: Write + WriteColor>(fmt: &mut W, basic_block: &BasicBlock, program: &Program) -> Result<()> {
+    for x in basic_block.mnemonics.iter() {
+        if !x.opcode.starts_with("__") {
+            display_mnemonic(fmt, &x, program)?;
         }
-    }).collect::<Vec<_>>();
-    display.iter().fold(seed, |acc, ref m| -> String {
-        format!("{}\n{}", acc, display_mnemonic(&m, program))
-    })
+    }
+    Ok(())
 }
 
 /// Displays the mnemonic in human readable form, and looks up any functions calls in `program`
-pub fn display_mnemonic(mnemonic: &Mnemonic, program: &Program) -> String {
+pub fn display_mnemonic<W: Write + WriteColor>(fmt: &mut W, mnemonic: &Mnemonic, program: &Program) -> Result<()> {
     let mut ops = mnemonic.operands.iter();
-    let mut fmt = format!( "{:8x}: {} ", mnemonic.area.start, mnemonic.opcode.bold().blue());
+    write!(fmt, "{:8x}: ", mnemonic.area.start)?;
+    color_bold!(fmt, Blue, mnemonic.opcode);
+    write!(fmt, " ")?;
     for token in &mnemonic.format_string {
         match token {
             &MnemonicFormatToken::Literal(ref s) => {
-                fmt = format!( "{}{}", fmt, format!("{}", s).bold().green());
+                color_bold!(fmt, Green, s);
             },
             &MnemonicFormatToken::Variable{ ref has_sign } => {
                 match ops.next() {
@@ -50,16 +69,16 @@ pub fn display_mnemonic(mnemonic: &Mnemonic, program: &Program) -> String {
                             } else { c };
                         let sign_bit = if s < 64 { 1u64 << (s - 1) } else { 0x8000000000000000 };
                         if !has_sign || val & sign_bit == 0 {
-                            fmt = format!( "{}{}", fmt, format!("{:x}",val).red());
+                            color!(fmt, Red, format!("{:x}", val));
                         } else {
-                            fmt = format!( "{}{:x}", fmt, (val as i64).wrapping_neg());
+                            color!(fmt, White, format!("{:x}", (val as i64).wrapping_neg()));
                         }
                     },
                     Some(&Rvalue::Variable{ ref name, subscript: Some(ref _subscript),.. }) => {
-                        fmt = format!( "{}{}", fmt, &name.to_lowercase().bold().white());
+                        color_bold!(fmt, White, &name.to_lowercase());
                     },
                     _ => {
-                        fmt = format!( "{}?", fmt);
+                        color!(fmt, Black, "?");
                     }
                 }
             },
@@ -71,32 +90,35 @@ pub fn display_mnemonic(mnemonic: &Mnemonic, program: &Program) -> String {
                                 let res = 1u64 << s;
                                 c % res
                             } else { c };
-                        let display = if is_code {
+                        if is_code {
                             if let Some(function) = program.find_function_by(|f| { f.start() == val }) {
-                                format!("{}{} <{}>", fmt, format!("{:x}",val).red(), function.name.yellow().bold())
+                                color!(fmt, Red, format!("{:x}",val));
+                                write!(fmt, " <", )?;
+                                color_bold!(fmt, Yellow, function.name);
+                                write!(fmt, ">")?;
                             } else {
-                                format!("{}{}", fmt, format!("{:x}",val).bold().purple())
+                                color_bold!(fmt, Magenta, format!("{:x}",val));
                             }
                         } else {
-                            format!("{}{:#x}", fmt, val)
-                        };
-                        fmt = display;
+                            write!(fmt, "{}", format!("{:#x}",val))?;
+                        }
                     },
                     Some(&Rvalue::Variable{ ref name, subscript: Some(_),.. }) => {
-                        fmt = format!( "{}{}", fmt, name.to_lowercase().yellow());
+                        color!(fmt, Yellow, name.to_lowercase());
                     },
                     Some(&Rvalue::Variable{ ref name, .. }) => {
-                        fmt = format!( "{}{}", fmt, name.to_lowercase().yellow());
+                        color!(fmt, Yellow, name.to_lowercase());
                     },
                     Some(&Rvalue::Undefined) => {
-                        fmt = format!( "{}{}", fmt, "undefined".bold().red());
+                        color_bold!(fmt, Red, "undefined");
                     },
                     None => {
-                        fmt = format!( "{}{}", fmt, "?".black().dimmed());
+                        color!(fmt, Black, "?");
                     }
                 }
             }
         }
     }
-    fmt
+    writeln!(fmt)?;
+    Ok(())
 }

cli/src/main.rs

@@ -12,7 +12,8 @@ extern crate futures;
 #[macro_use]
 extern crate log;
 extern crate env_logger;
-extern crate colored;
+extern crate termcolor;
+extern crate atty;
 
 use panopticon_amd64 as amd64;
 use panopticon_analysis::analyze;
@@ -37,6 +38,8 @@ use errors::*;
 #[derive(StructOpt, Debug)]
 #[structopt(name = "panop", about = "A libre cross-platform disassembler.")]
 struct Args {
+    #[structopt(long = "color", help = "Always color")]
+    color: bool,
     /// Print every function the function calls
     #[structopt(short = "c", long = "calls", help = "Print every address of every function this function calls")]
     calls: bool,
@@ -69,7 +72,7 @@ impl Filter {
     }
     pub fn is_match(&self, func: &Function) -> bool {
         if let Some(ref name) = self.name {
-            if name == &func.name { return true }
+            if name == &func.name || func.aliases().contains(name){ return true }
         }
         if let Some(ref addr) = self.addr {
             return *addr == func.start()
@@ -102,7 +105,7 @@ fn disassemble(args: Args) -> Result<()> {
     });
 
     for function in functions {
-        println!("{}", display::display_function(&function, &program));
+        display::print_function(&function, &program, args.color)?;
         if args.calls {
             let calls = function.collect_call_addresses();
             println!("Calls ({}):", calls.len());