#!/usr/bin/perl
# usage: 
# nema.pl inputfile password

use Term::ANSIColor;
($infile, $password) = @ARGV;

print "NeMa cipher simulation 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);
@rwE = ( 5, 18, 20, 23, 12,  1, 22, 10, 14,  3, 17,  0, 24,
	21,  7,  6,  2, 19, 25,  9,  8, 11, 16, 13, 15,  4);
@rwF = (14,  3,  5, 10, 25,  7,  6, 18, 23, 16,  0, 15,  2,
	19, 21,  9,  1, 20, 17, 13, 24, 22,  8,  4, 12, 11);
# reflector
@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
@irwF = (10, 16, 12,  1, 23,  2,  6,  5, 22, 15,  3, 25, 24, 
	 19,  0, 11,  9, 18,  7, 13, 17, 14, 21,  8, 20,  4);
@irwE = (11,  5, 16,  9, 25,  0, 15, 14, 20, 19,  7, 21,  4, 
	 23,  8, 24, 22, 10,  1, 17,  2, 13,  6,  3, 12, 18);
@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);

# 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);
@nr17 = (0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0,
	 0, 0, 0, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0);
@nr18 = (1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
	 0, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 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);

# inner key: 4 rotor wheels and 6 notch rings, from right to left!
# eg. 12A-13B-14C-15D-22/1
my @wheels    = (\@rwD,  \@rwC,  \@rwB,  \@rwA);  # contact wheels
my @invWheels = (\@irwD, \@irwC, \@irwB, \@irwA);
my @rings = (\@nr1, \@nr22, \@nr15, \@nr14, \@nr13, \@nr12); # notch rings

# outer key (10-letter password)
@state = &getInitialState($password);
print "initial state " . &displayState() . "\n";

&checkForFile($infile);
$text = &stripText(`cat $infile`);    
@clearText = split('', $text);
@cipherText = ();
$nSteps = 0;
@wheelFreq  = (0) x 10;
@letterFreq = (0) x 26;

# loop on message letters
foreach $letter (@clearText) {
    &encrypt($letter);
}
print "\nciphertext:\n";
&blockPrint(@cipherText);

# &printWheelFreq();
# &printLetterFreq();

sub getInitialState {
    my $password = shift;
    $password = lc($password);
    die "password $password is too short!" if (length($password)<10);
    my @initState = ();
    my @letters = split '', $password;
    for ($i=9; $i>=0; $i--) {
	push @initState, ((ord($letters[$i]) - ord('a') +18) % 26);
    }
    return @initState;
}

sub stepState {
    my @oldState = @state;

    # test current notch ring states
    my $nrState0 = $rings[1]->[($state[0]+24) % 26];
    my $nrState1 = $rings[0]->[($state[0]+25) % 26];
    my $nrState2 = $rings[2]->[($state[2]+24) % 26];
    my $nrState4 = $rings[3]->[($state[4]+24) % 26];
    my $nrState6 = $rings[4]->[($state[6]+24) % 26];
    my $nrState8 = $rings[5]->[($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 ($nrState0==1); # wheel 2 rotates if left notch ring of drive wheel 1 is active ( = 1)
    $state[5] = ($state[5]-1) % 26 if ($nrState4==1); # wheel 6 rotates if notch ring of drive wheel 5 is active
    $state[9] = ($state[9]-1) % 26 if ($nrState8==1); # wheel 10 rotates if notch ring of drive wheel 9 is active
    if ($nrState1==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 ($nrState2==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 ($nrState6==1); # wheel 8 rotates if notch ring of drive wheel 1 is active 
                                                          # AND notch ring of drive wheel 7 is active
    }
    for ($i=0; $i<10; $i++) {
	 $wheelFreq[$i]++ if ($state[$i] != $oldState[$i]);
    }
}

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

sub encrypt {
    $nSteps++;
    printf "step %2i ", $nSteps;
    &stepState();
    print "state " . &displayState();

    my $clearChar = shift;

    my $channel = $keybd[&encode($clearChar)];
    $channel = ($channel-2) % 26;
    print " cleartext $clearChar ";
    #printf " channel %2i ", $channel;
    # electric current from right to left
    for ($i=0; $i<4; $i++) {
	my $offset = &getWheelOffset($i);
	$channel = $wheels[$i]->[($channel+$offset) % 26];
	#if ($i==0) {printf "offED %2i outD %2i ", $offset, $channel};
	#if ($i==1) {printf "offDC %2i outC %2i ", $offset, $channel};
	#if ($i==2) {printf "offCB %2i outB %2i ", $offset, $channel};
	#if ($i==3) {printf "offBA %2i outA %2i ", $offset, $channel};
    }

    # electric current through reflector
    my $offset = &getWheelOffset(4);
    #printf "ukw-in %2i ", ($channel+$offset) % 26;
    $channel = $ukw[($channel+$offset) % 26];
    #printf "ukw-out %2i ", $channel;
    
    # electric current from left to right
    for ($i=3; $i>=0; $i--) {
	my $offset = &getWheelOffset($i+1);
	#$channel = &reverseWheel(($channel-$offset) % 26, $i);
	$channel = $invWheels[$i]->[($channel-$offset) % 26];
	#if ($i==3) {printf "offUA %2i outA %2i ", $offset, $channel};
    }

    # offset to entry plate
    $offset = &getWheelOffset(0);
    $channel = ($channel-$offset) % 26;
    my $cipherChar = $lamps[(26-$channel-2) % 26];
    #printf "lamp-in %2i ", $channel;
    print "cipher $cipherChar\n";
    push @cipherText, $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]) % 26;
}

sub checkText {
# Remove everything that is not in the alphabet
    my($input,$text,$test,$i,@tab);
    $input=$_[0];
    chomp($input);
    $text="";
   
    @tab=split('',$input);
    foreach $i (@tab) {
        $test=ord($i);
        if(($test >= ord('A') && $test <= ord('Z')) || ($test >= ord('a') && $test <= ord('z'))) {
	    $text .= $i;
        }
    }
    return($text);   
}

sub stripText {
    # remove whitespace and punctuation, convert to lower-case
    my($text);
    while ($text = shift) {
	chomp $text;
	$text = &checkText($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 printWheelFreq {
    print "\nWheel frequencies:\n";
    for ($i=9; $i>=0; $i--) {
	printf "%1.3f ", $wheelFreq[$i]/$nSteps; 
    }
    print "\n\n";
}

sub printLetterFreq {
    print "\nLetter frequencies:\n";
    for ($i=0; $i<26; $i++) {
	my $letter = &decode($i);
	printf "%2i $letter  %1.3f\n", $i, $letterFreq[$i]/$nSteps; 
    }
}

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);
}