#!/usr/bin/perl
# usage: 
# known-plaintext.pl cipher-file plain-file

use Term::ANSIColor;
($cipherfile, $plainfile) = @ARGV;

print "NeMa known-plaintext attack by nicZ && chE\n";

# keyboard and lampboard
@lamps = qw/q w e r t z u i o p
             a s d f g h j k l
              y x c v b n m/;
# electric channel connections to keyboard letters
@keybd = (16, 3, 5, 14, 24, 13, 12, 11, 19, 10, 9, 8, 1,
	  2, 18, 17, 0, 23, 15, 22, 20, 4, 25, 6, 7, 21);

# rotor wheel connections
@rwA = ( 4, 13, 14, 18, 12,  1, 21,  9,  3, 17, 15, 25, 23, 
	 8, 22, 24,  7, 19,  5, 10,  2,  0, 11, 20,  6, 16);
@rwB = ( 3,  6, 17,  8, 19, 14,  7, 10, 15,  0,  9, 23, 18,
	24, 12, 21, 13, 20,  2,  1, 16,  5, 11,  4, 22, 25);
@rwC = (17, 16, 18, 25, 13,  9, 14,  6,  1, 10, 24, 19,  8, 
	 4,  2, 20, 15,  3,  7,  5, 23, 12, 11, 22, 21,  0);
@rwD = ( 4, 21, 11, 18, 17,  2, 15,  7,  0,  8, 25, 10, 13, 
	 3, 24,  6,  5,  9, 14,  1, 19, 22, 16, 23, 12, 20);
@ukw = (13, 11, 18, 17, 12, 21, 16, 20, 15, 10,  9,  1,  4,
	 0, 19,  8,  6,  3,  2, 14,  7,  5, 24, 25, 22, 23);
# rotor wheel connections backwards
@irwD = ( 8, 19,  5, 13,  0, 16, 15,  7,  9, 17, 11,  2, 24, 
	 12, 18,  6, 22,  4,  3, 20, 25,  1, 21, 23, 14, 10);
@irwC = (25,  8, 14, 17, 13, 19,  7, 18, 12,  5,  9, 22, 21, 
	  4,  6, 16,  1,  0,  2, 11, 15, 24, 23, 20, 10,  3);
@irwB = ( 9, 19, 18,  0, 23, 21,  1,  6,  3, 10,  7, 22, 14, 
	 16,  5,  8, 20,  2, 12,  4, 17, 15, 24, 11, 13, 25);
@irwA = (21,  5, 20,  8,  0, 18, 24, 16, 13,  7, 19, 22,  4,
	  1,  2, 10, 25,  9,  3, 17, 23,  6, 14, 12, 15, 11);
my @wheels = (\@rwD, \@rwC, \@rwB, \@rwA);
my @invWheels = (\@irwD, \@irwC, \@irwB, \@irwA);

# notch rings
@nr1  = (0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 
	 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 1);
@nr12 = (0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0,
	 0, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1);
@nr13 = (1, 1, 0, 1, 1, 1, 1, 0, 0, 1, 1, 0, 1, 
	 1, 0, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 0);
@nr14 = (0, 0, 1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1,
	 1, 1, 1, 1, 0, 1, 0, 0, 1, 0, 1, 0, 1);
@nr15 = (1, 0, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 1, 
	 0, 1, 1, 1, 1, 1, 1, 0, 1, 0, 1, 1, 1);
@nr22 = (1, 1, 0, 0, 1, 0, 1, 1, 0, 0, 1, 0, 1,
	 1, 0, 1, 1, 1, 1, 0, 0, 1, 1, 1, 0, 0);

# get ciphertext and plaintext
&checkForFile($cipherfile);
&checkForFile($plainfile);
@cipherText = split('', &stripText(`cat $cipherfile`));
@plainText  = split('', &stripText(`cat $plainfile`));
my $cipherLength = scalar(@cipherText);
my $plainLength  = scalar(@plainText);
print "ciphertext length: $cipherLength chars, plaintext length: $plainLength chars\n";
warn " ciphertext and plaintext have different lengths!\n" if ($cipherLength != $plainLength);

print "\nplaintext:\n";  &blockPrint(@plainText);
print "\nciphertext:\n"; &blockPrint(@cipherText);

$nSteps = 0;
@state = ();
$nKeys = 0;
@decodedText = ();
$startTime = time();

# loop on possible passwords
for ($pwd9=18; $pwd9<44; $pwd9++) {
    for ($pwd8=18; $pwd8<44; $pwd8++) {
	for ($pwd7=18; $pwd7<44; $pwd7++) {
	    for ($pwd6=18; $pwd6<44; $pwd6++) {
		for ($pwd5=18; $pwd5<44; $pwd5++) {
		    for ($pwd4=18; $pwd4<44; $pwd4++) {
			for ($pwd3=18; $pwd3<44; $pwd3++) {
			    for ($pwd2=18; $pwd2<44; $pwd2++) {
				for ($pwd1=18; $pwd1<44; $pwd1++) {
				    for ($pwd0=18; $pwd0<44; $pwd0++) {
					@state = ($pwd0, $pwd1, $pwd2, $pwd3, $pwd4, 
						  $pwd5, $pwd6, $pwd7, $pwd8, $pwd9);
					foreach (@state) {$_ %= 26};
					$nKeys++;
					my $password = &displayState();
					print "trying password $nKeys: $password  rate: " 
					    . &getRate() ."\r" if (($nKeys % 100)==0);
					@decodedText = ();
					$found = 1;
					# loop on message letters
					for ($k=0; $k<$cipherLength; $k++) {
					    &encrypt($cipherText[$k]);
					    $found=0 if ($decodedText[$k] ne $plainText[$k]);
					    last if ($found==0);
					}
					next if ($found==0);
					print "\n\ndecoded text:\n"; &blockPrint(@decodedText); 
					print " key found: "; print color 'green'; 
					print "$password\n";  print color 'reset';
					exit;
				    }
				}
			    }
			}
		    }
		}
	    }
	}
    }
}

sub getRate {
    return "" if (time() == $startTime);
    $rate = int($nKeys/(time() - $startTime));
    return "$rate keys/second";
}

sub stepState {
    # save current notch ring states
    my $state0 = ($state[0]+24) % 26;
    my $state1 = ($state[0]+25) % 26;
    my $state2 = ($state[2]+24) % 26;
    my $state4 = ($state[4]+24) % 26;
    my $state6 = ($state[6]+24) % 26;
    my $state8 = ($state[8]+24) % 26;

    # rotate wheels
    $state[0] = ($state[0]-1) % 26; # Red wheel rotates for each step
    $state[4] = ($state[4]-1) % 26; # wheel 5 rotates for each step
    $state[8] = ($state[8]-1) % 26; # wheel 9 rotates for each step
    $state[1] = ($state[1]-1) % 26 if ($nr22[$state0]==1); # wheel 2 rotates if left notch ring of drive wheel 1 is active ( = 1)
    $state[5] = ($state[5]-1) % 26 if ($nr14[$state4]==1); # wheel 6 rotates if notch ring of drive wheel 5 is active
    $state[9] = ($state[9]-1) % 26 if ($nr12[$state8]==1); # wheel 10 rotates if notch ring of drive wheel 9 is active
    if ($nr1[$state1]==1) {
	$state[2] = ($state[2]-1) % 26; # wheel 3 rotates if right notch ring of drive wheel 1 is active
	$state[6] = ($state[6]-1) % 26; # wheel 7 rotates if right notch ring of drive wheel 1 is active
	$state[3] = ($state[3]-1) % 26 if ($nr15[$state2]==1); # wheel 4 rotates if notch ring of drive wheel 1 is active 
	                                                       # AND notch ring of drive wheel 3 is active
	$state[7] = ($state[7]-1) % 26 if ($nr13[$state6]==1); # wheel 8 rotates if notch ring of drive wheel 1 is active 
                                                               # AND notch ring of drive wheel 7 is active
    }
}

sub displayState {
    my $strState = "";
    for ($i=9; $i>=0; $i--) {
	$strState .= chr(($state[$i]-18) % 26 + ord('A'));
    }
    return $strState;
}

sub encrypt {
    $nSteps++;
    &stepState();
    my $clearChar = shift;

    my $channel = $keybd[&encode($clearChar)];
    $channel = ($channel-2) % 26;
    # electric current from right to left
    for ($i=0; $i<4; $i++) {
	my $offset = &getWheelOffset($i);
	$channel = $wheels[$i]->[($channel+$offset) % 26];
    }

    # electric current through reflector
    $channel = $ukw[($channel+&getWheelOffset(4)) % 26];
    
    # electric current from left to right
    for ($i=3; $i>=0; $i--) {
	my $offset = &getWheelOffset($i+1);
	$channel = $invWheels[$i]->[($channel-$offset) % 26];
    }

    # offset to entry plate
    $channel = ($channel-&getWheelOffset(0)) % 26;
    my $cipherChar = $lamps[(26-$channel-2) % 26];

    # save letter frequencies
    push @decodedText, $cipherChar;
    $letterFreq[&encode($cipherChar)]++;
}

sub getWheelOffset {
    # get offset between to adjacent contact wheels
    my $wheelIndex = shift;
    return $state[1] if $wheelIndex == 0;
    return ($state[2*$wheelIndex+1] - $state[2*$wheelIndex-1]);
}

sub stripText {
    # remove whitespace and punctuation, convert to lower-case
    my $stext = "";
    while ($text = shift) {
	chomp $text;
	$stext .= $text;
    }
    $stext =~ s/( |\.|,|\n)//g;
    return lc($stext);
}

sub encode {
    # convert [a-z] to [0-25]
    $char = shift;
    return (ord($char) - ord('a'));
}

sub decode {
    # convert [0-25] to [a-z]
    $code = shift;
    return chr($code + ord('a'));
}

sub blockPrint {
    # print text in blocks of five upper-case letters
    my $i = 0;
    while ($char = shift) {
	print " "  if ($i %  5 == 0 && $i > 0);
	print "\n" if ($i % 40 == 0 && $i > 0);
	print uc($char);
	$i++;
    }
    print "\n";
}

sub checkForFile {
    my ($fn) = @_;
    die "Error : file $fn does not exist!\n" unless (-e $fn);
    warn "Warning : file $fn has zero size!\n" if (-z $fn);
}