use Tk;
use Tk::ROTextANSIColor;
use Term::ANSIColor qw(:constants);
#use GD::Graph::lines;
use strict;
no strict 'subs';

my %disasm;
my $w = 600;
my $h = 400;
my $searchdistance = 10;
my $searchsize = 5;
my $max = OllyDbg::MAXCMDSIZE;
my $da = OllyDbg::DISASM_ALL;
my $graphnum = 0;
my (@d1, @d2, @d3, @d4, @a1, @a2);
my (@asm1, @asm2);
my ($col1, $col2);
my ($end, $d1end, $d2end);
my $text;
#my $red = WHITE . ON_RED;
my $red = RED;
my $black = RESET;
my $blue = BLUE;

my @blank;
for (0..$searchsize-1) {
  push(@blank, undef);
}

my $g1 = "diff1.png";
my $g2 = "diff2.png";
my $g3 = "diffresult1.png";
my $g4 = "diffresult2.png";

my $exe1 = "";
my $exe2 = "";
my $file1open = 0;
my $file2open = 0;
my $loopcount = 0;
my $stat;
sub mainLoop {
  $stat = OllyDbg::Getstatus();
  if ($file1open == 0) 
  {
    OllyDbg::Browsefilename("Select first file for diff", $exe1, ".exe;*.dll", 0);
    OllyDbg::OpenEXEfile($exe1,0);
    $file1open = 1;
  } 
  elsif (($$stat == OllyDbg::STAT_STOPPED) && ($file1open == 1) && ($file2open == 0))
  {
    createGraph($g1,$exe1);
    OllyDbg::Browsefilename("Select second file for diff", $exe2, ".exe;*.dll", 0);
    OllyDbg::OpenEXEfile($exe2,0);
    $file2open = 1;
  }
  elsif (($$stat == OllyDbg::STAT_STOPPED) && ($file2open == 1))
  {
    createGraph($g2,$exe2);
    $text = sprintf("%s%s%32s%s  %s%-32s%s\n", " "x8, 	
	RED, $exe1, RESET,
	RED, $exe2, RESET);

    $d1end = scalar(@d1) - 1;
    $d2end = scalar(@d2) - 1;
    $end = ($d2end > $d1end) ? $d2end: $d1end;

    calculateDifference($g3,$g4);

    open (OUT, ">C:\\diff.ansi");
    print OUT $text;
    close OUT;

    my $mw = new MainWindow('-title' => 'Phase Cancellation Analysis ASM Listing');
    my $body = $mw->Frame->pack(-anchor=>'center', -expand=>'yes', -fill=>'both');

    my $t = $body->Scrolled(qw/ROTextANSIColor -relief sunken
         -height 40
         -scrollbars e
         -width 88/);
    $t->pack(-expand=>'yes', -fill=>'both');
    $t->insert('end', $text);

    MainLoop();
    exit;

  }
}

sub createGraph {
  my $offset = 0;
  my $buf;
  my (@data, @addr);
  my $graphfile = shift;
  my $exe = shift;

  OllyDbg::Addtolist(0,0,"Opening %s",$exe);
  open(IN,$exe) or die "Couldn't open $exe : $!\n";

  read(IN, $buf, 2);
  die "DOS header not found!\n" unless $buf eq "MZ";
  seek(IN, 60, 0);
  read(IN, $buf, 4);
  my $peoffset = unpack("I",$buf);

  seek(IN, $peoffset, 0);
  read(IN, $buf, 2);
  die "PE header not found!\n" unless $buf eq "PE";
  OllyDbg::Addtolist(0,-1,"PE header found at offset %x", $peoffset);

  seek(IN, $peoffset+20, 0);
  read(IN, $buf, 2);
  my $optheadersize = unpack("S", $buf);
  OllyDbg::Addtolist(0,-1,"Optional header size = %x", $optheadersize);
  
  seek(IN, $peoffset+28, 0);
  read(IN, $buf, 4);
  my $codesize = unpack("I", $buf);
  OllyDbg::Addtolist(0,-1,"Code size = %x", $codesize);
  
  seek(IN, $peoffset+44, 0);
  read(IN, $buf, 4);
  my $codebase = unpack("I", $buf);
  OllyDbg::Addtolist(0,-1,"Code base = %x", $codebase);

  seek(IN, $peoffset+52, 0);
  read(IN, $buf, 4);
  my $imagebase = unpack("I", $buf);
  OllyDbg::Addtolist(0,-1,"Image base = %x", $imagebase);

  seek(IN, $peoffset+24+$optheadersize+12, 0);
  read(IN, $buf, 4);
  my $codeva = unpack("I", $buf);
  OllyDbg::Addtolist(0,-1,"Code VA = %x", $codeva);

  die "Code segment not first (code may be packed?)" if $codeva != $codebase;
  
  seek(IN, $peoffset+24+$optheadersize+16, 0);
  read(IN, $buf, 4);
  my $coderawsize = unpack("I", $buf);
  OllyDbg::Addtolist(0,-1,"Code raw size = %x", $coderawsize);

  seek(IN, $peoffset+24+$optheadersize+20, 0);
  read(IN, $buf, 4);
  my $coderawpointer = unpack("I", $buf);
  
  seek(IN, $coderawpointer, 0);
  read(IN, $buf, $coderawsize);
  close IN;

  my $start = $imagebase + $codebase;
  my $stop = $start + $coderawsize;

  $graphnum++;
  while ($offset < $coderawsize) {
    my $asm;
    my $cmd = substr($buf,$offset,$max);
    my $psize;
    my $decode = OllyDbg::Finddecode($start+$offset,$psize);
    my $ret = OllyDbg::Disasm($cmd,$max,$start+$offset,$decode,\%disasm,$da,NULL);
    if ($ret) {
      $asm = $disasm{result};
    } else {
      $asm = "(BAD)";
    }
    if ($graphnum == 1) {
      push(@asm1, $asm);
    } else {
      push(@asm2, $asm);
    }
    my $promille = ($offset / $coderawsize) * 1000;
    OllyDbg::Progress($promille, "Translating ASM to wave (%s)", $asm);
    my $bval;
    if ($asm =~ /PUSH EBP/i) {
      $bval = 255;
    } elsif ($asm =~ /^RET/i) {
      $bval = 200;
    } elsif ($asm =~ /^INT/i) {
      $bval = 210;
    } elsif ($asm =~ /LEAVE/i) {
      $bval = 240;
    } elsif ($asm =~ /^JMP/i) {
      $bval = 45;
    } elsif ($asm =~ /^JE/i) {
      $bval = 40;
    } elsif ($asm =~ /^JN/i) {
      $bval = 35;
    } elsif ($asm =~ /^JB/i) {
      $bval = 30;
    } elsif ($asm =~ /^J/i) {
      $bval = 25;
    } elsif ($asm =~ /^CALL/i) {
      if ($asm =~ /JMP/i) {
        $bval = 175;
      } elsif($asm =~ /DWORD/i) {
        $bval = 160;
      } elsif($asm =~ /CALL E[A-Z][XIP]/i) {
        $bval = 150;
      } else {
        $bval = 50;
      }
    } elsif ($asm =~ /^CMP/i) {
      $bval = 20;
    } elsif ($asm =~ /^TEST/i) {
      $bval = 18;
    } elsif ($asm =~ /^MOV/i) {
      $bval = 15;
    } elsif ($asm =~ /^LEA/i) {
      $bval = 10;
    } elsif ($asm =~ /^XOR/i) {
      $bval = 5;
    } elsif ($asm =~ /^PUSH/i) {
      $bval = 8;
    } elsif ($asm =~ /^POP/i) {
      $bval = 3;
    } else { $bval = 0; }
    if ($graphnum == 2) { $bval = -$bval }
    push(@data, $bval);
    push(@addr, sprintf("%x", $start+$offset));
    if ($graphnum == 1) {
      push(@d1, $bval);
      push(@a1, sprintf("%x", $start+$offset));
    } elsif ($graphnum == 2) {
      push(@d2, $bval);
      push(@a2, sprintf("%x", $start+$offset));
    }
    if ($ret) {
      $offset += $ret;
    } else {
      $offset++;
    }
  }
  OllyDbg::Progress(1000, "Translating ASM to wave");
  OllyDbg::Progress(0, "");
  my $ymax = 300;
  my $title = "Wave plot for $exe1";
  if ($graphnum == 2) {
    $title = "Inverted wave plot for $exe2";
    $ymax = 0;
  }
  return;
  my @graphdata = ([@addr],[@data]);
 
  my $graph = GD::Graph::lines->new($w, $h/4);
           $graph->set(
               x_label           => 'Address',
               y_label           => 'ASM Translation Value',
               x_label_skip	=>  $coderawsize / 5,	
               title             => $title,
               y_max_value       => $ymax,
  	     bgclr		=> '#FFFFFF',
  	     fgclr		=> '#0000FF',
  	     labelclr		=> '#FF0000',
	     dclrs 		=> [ qw(lgreen red) ],
  	     line_width		=> 1,
  	     transparent	=> 0
           ) or die $graph->error;
   my $gd = $graph->plot(\@graphdata) or do {
    OllyDbg::Error($graph->error);
    die $!
    };
  
  open(IMG, ">$graphfile") or die $!;
  binmode IMG;
  print IMG $gd->png;
  close IMG;
}
  

sub calculateDifference {
  my ($graphfile1,$graphfile2) = @_;

  my $o1 = 0;
  my $o2 = 0;
  my $diff;
  my ($searchoffset, $total); 
  
DIFFLOOP: for (0..$end) {
    OllyDbg::Progress(($_ / $end) * 1000, "Calculating difference");
    $diff = $d1[$o1] + $d2[$o2];
    if ($diff == 0) {
      push(@d3,undef); push(@d4,undef);
      # print matched instructions
#      $text .= sprintf("%s%s%s  %32s  %-32s  %s%s%s\n", 
#	$blue, $a1[$o1], $black,
#	$asm1[$o1], $asm2[$o2],
#	$blue, $a2[$o2], $black);
    } else {
      #  we hit a different instruction.  Try to resync
      # grab search chunk to find in other binary
      my @searcher1 = @d1[$o1..$o1+$searchsize-1];
      my @searcher2 = @d2[$o2..$o2+$searchsize-1];
      
      OllyDbg::Progress(($_ / $end) * 1000, "Trying to resync");      

      my $searchend = ($o2+$searchdistance > $d2end) ? $d2end : $o2+$searchdistance;
      for $searchoffset ($o2+1..$searchend) {
        my @searchin = @d2[$searchoffset..$searchoffset+$searchsize-1];
        $total = sumArrays(\@searcher1, \@searchin);
	if ($total == 0) {
	  # sync found, copy d2 changed chunk we just skipped to output array
	  push(@d4, @d2[$o2..$searchoffset-1]);
	  # print d2 difference
	  my $spacelen = length($a1[$o1]);
	  for(0..$searchoffset-1-$o2) {
            $text .= sprintf("%s%s%s  %32s  %s%-32s%s  %s%s%s\n", 
	    $blue," "x $spacelen,$black, 
	    " "x32,$red,$asm2[$o2+$_],$black, 
	    $blue,$a2[$o2+$_],$black);
	  }
	  # skip to end of known matching chunk;
	  push(@d3,@blank);
	  push(@d4,@blank);
	  # print matched block
	  for (0..$searchsize) { 
	    $text .= sprintf("%s%s%s  %s%32s  %-32s  %s%s%s\n",
	    $blue,$a1[$o1+$_],$black,
	    $black,$asm1[$o1+$_],$asm2[$searchoffset+$_],
	    $blue,$a2[$searchoffset+$_],$black);
	  }
	  $o2 = $searchoffset + $searchsize - 1;
	  $o1 += $searchsize - 1;
	  goto ENDSEARCH;
	}
      }
      $searchend = ($o1+$searchdistance > $d1end) ? $d1end : $o1+$searchdistance;
      for $searchoffset ($o1+1..$searchend) {
        my @searchin = @d1[$searchoffset..$searchoffset+$searchsize-1];
        $total = sumArrays(\@searcher2, \@searchin);
	if ($total == 0) {
	  # sync found, copy d1 changed chunk we just skipped to output array
	  push(@d3, @d1[$o1..$searchoffset-1]);
	  # print d1 difference
	  for(0..$searchoffset-1-$o1) {
            $text .= sprintf("%s%s%s  %32s%s\n", 
	    $blue,$a1[$o1+$_],$red, $asm1[$o1+$_],$black); 
	  }
	  # skip to end of known matching chunk on d2;
	  push(@d3,@blank);
	  push(@d4,@blank);
	  # print matched block
	  for (0..$searchsize) { 
	    $text .= sprintf("%s%s%s  %s%32s  %-32s  %s%s%s\n",
	    $blue,$a1[$searchoffset+$_],$black,
	    $black,$asm1[$searchoffset+$_],$asm2[$o2+$_],
	    $blue,$a2[$o2+$_],$black
	    );
	  }
	  $o1 = $searchoffset + $searchsize - 1;
	  $o2 += $searchsize - 1;
	  goto ENDSEARCH;
	}
      }
      # no match. just log instruction into d3 and d4 and move on
      push(@d3, $d1[$o1]);
      push(@d4, $d2[$o2]);
      
      #print unmatched instruction
      $text .= sprintf("%s%s%s  %s%32s  %-32s%s  %s%s%s\n", 
      $blue, $a1[$o1], $black,
      $red, $asm1[$o1], $asm2[$o2], $black,
      $blue, $a2[$o2], $black);
ENDSEARCH:
    } # end of if diff == 0 else;
    $o1++; $o2++;
  }
  OllyDbg::Progress(0, "");
  return;
  my @graphdata1 = ([@a1],[@d1],[@d3]);
  my @graphdata2 = ([@a2],[@d2],[@d4]);
  my $size = scalar(@d1);
  my $graph1 = GD::Graph::lines->new($w, $h/4);
           $graph1->set(
               x_label           => 'Address',
               y_label           => 'ASM Translation Value',
        	     x_label_skip	=> $size / 5,
               title             => 'Difference Plot 1',
               y_min_value       => 0,
               y_max_value       => 300,
  	     bgclr		=> '#FFFFFF',
  	     fgclr		=> '#0000FF',
  	     labelclr		=> '#FF0000',
	     skip_undef		=> 1,
	     dclrs 		=> [ qw(lgreen red) ],
  	     line_width		=> 1,
  	     transparent	=> 0
           ) or die $graph1->error;
  my $gd1 = $graph1->plot(\@graphdata1) or do {
    OllyDbg::Error($graph1->error);
    die $!
    };
  my $graph2 = GD::Graph::lines->new($w, $h/4);
           $graph2->set(
               x_label           => 'Address',
               y_label           => 'ASM Translation Value',
        	     x_label_skip	=> $size / 5,
               title             => 'Difference Plot 2',
               y_min_value       => -300,
               y_max_value       => 0,
  	     bgclr		=> '#FFFFFF',
  	     fgclr		=> '#0000FF',
  	     labelclr		=> '#FF0000',
	     skip_undef		=> 1,
	     dclrs 		=> [ qw(lgreen red) ],
  	     line_width		=> 1,
  	     transparent	=> 0
           ) or die $graph2->error;
   my $gd2 = $graph2->plot(\@graphdata2) or do {
    OllyDbg::Error($graph2->error);
    die $!
    };
  
  open(IMG, ">$graphfile1") or die $!;
  binmode IMG;
  print IMG $gd1->png;
  close IMG;

  open(IMG, ">$graphfile2") or die $!;
  binmode IMG;
  print IMG $gd2->png;
  close IMG;
}

sub sumArrays {
  my ($a1, $a2) = @_;
  return -1 if (scalar(@{$a1}) != scalar (@{$a2}));
  my $last = scalar(@{$a1}) - 1;
  for (0..$last) {
    return -1 if $$a1[$_] + $$a2[$_] != 0; # found difference
  }
  return 0; # no difference in these chunks
}