#!/usr/bin/perl -w
use strict;
use Fcntl;
use IO::Socket::INET;

# There is a sync bug which needs to be fixed
use constant HOSTNAME=>'b00b.ee.cooper.edu'; # sorry about this, the other guys named the computer :P
use constant DEBUG=>0;

################################################################################################
#PERSONAL IDENTY 
{
	package iden;
	use fields qw(name sk pk shared_key hmac_key _dh _iv);
	use Crypt::Rijndael;
	use Crypt::DH;
	use Crypt::Random qw( makerandom );
	use Digest::SHA qw(sha256 sha256_hex);
	use Digest::HMAC_SHA1 qw(hmac_sha1 hmac_sha1_hex);

	sub new {
		my iden $self=shift;
		unless (ref $self) { $self=fields::new($self); }
		$self->{name}=shift || die "please specify name: $!\n";
		$self->{hmac_key}=shift || die "please specify HMAC shared key\n";
		#DH related --------------------------------------------------------------------
		#100digit primes from http://primes.utm.edu/lists/small/small.html#100
		my $p=Math::BigInt->new('5915587277');
#		2074722246773485207821695222107608587480996474721117292752992589912196684750549658310084416732550077');
		my $g=Math::BigInt->new('2');
#		2367495770217142995264827948666809233066409497699870112003149352380375124855230068487109373226251983');

		my $r=Crypt::Random::makerandom( Size => 100, Strength => 0 ); 

		$self->{_dh}=Crypt::DH->new;
		($self->{_dh})->g($g);
		($self->{_dh})->p($p);
		($self->{_dh})->priv_key($r);
		($self->{_dh})->generate_keys;
		$self->{pk}=($self->{_dh})->pub_key;
		$self->{sk}=($self->{_dh})->priv_key;
		$self->{shared_key}=0;
		#DH related --------------------------------------------------------------------
		$self->{_iv}="pimpupthecryptoz"; #IV, should be random, but who cares for this proof of concept

		return $self;
	}
	sub hmac_rekey {
		 my iden $self=shift;
		 unless (ref $self) { die "please use rekey properly $!\n"; }
#		 $self->{hmac_key}="";
#		 $self->{hmac_key}.=$self->{shared_key}; #otherwise it copies the reference
		 $self->{hmac_key}=sha256_hex($self->{shared_key});
	}
	sub rekey {
		 my iden $self=shift;
		 unless (ref $self) { die "please use rekey properly $!\n"; }
		 my $r=Crypt::Random::makerandom( Size => 100, Strength => 0 );  
		 ($self->{_dh})->priv_key($r);
		 ($self->{_dh})->generate_keys;
		$self->{pk}=($self->{_dh})->pub_key;
		$self->{sk}=($self->{_dh})->priv_key;
	}
	sub exchange {
		 my iden $self=shift;
		 my $pub_key=shift || die "please specify the other pub key in exchange $!\n";
		 unless (ref $self) { die "please use exchange properly $!\n"; }
		 $self->{shared_key}=($self->{_dh})->compute_secret($pub_key);
	}
	sub print_info {
		my iden $self=shift;
		unless (ref $self) { die "please use print_info properly $!\n"; }
		print "name: ".$self->{name}."\npk=". $self->{pk}."\nsk=". $self->{sk}."\nshared_key=". $self->{shared_key}."\nhmac_key=".$self->{hmac_key}."\n\n";
	}
	#_padd(message) 0-pad messages to 16 byte alignment
	sub _padd($) {
		my $m=shift or die "no m supplied to _padd: $!";
		my $lm=length($m)%16;
		if($lm ne 0) {
#TODO: don't just zeropadd, add a 16byte filed that says how much you zeropad, but thats for later...
#			my $a=sprintf("%016d",16675-$lm);
#			print "a= ".(16-$lm)."=\'$a\'\n";
			return $m.("\0"x(16-$lm));
		} else {
			return $m;
		}
	}

	#sign(m) sign messageg with key=hmac_key
	sub sign {
		my iden $self=shift;
		unless (ref $self) { die "please use sign properly $!\n"; }
		my $m=shift || die "please specify a message to sign $!\n";
		return hmac_sha1_hex($m,$self->{hmac_key});
	}

	#verify(m1,m2) verify that sign(m2)=m1 with key=hmac_key
	sub verify {
		my iden $self=shift;
		unless (ref $self) { die "please use verify properly $!\n"; }
		my $m1=shift || die "please specify a message to verify $!\n";
		my $m2=shift || die "please specify a message to verify $!\n";
		return ($m1 eq $self->sign($m2));
	}

	#encrypt(message) encrypt message with key=sha256(shared_key)
	sub encrypt { 
		my iden $self=shift;
		unless (ref $self) { die "please use encrypt properly $!\n"; }
		my $plaintext=shift || die "please specify a plaintext to encrypt $!\n";
		my $cipher= Crypt::Rijndael->new(sha256($self->{shared_key}),Crypt::Rijndael::MODE_CTR());
		$cipher->set_iv($self->{_iv});
		my $ciphertext  = $cipher->encrypt(_padd($plaintext));
		return $ciphertext;
	}
	#decrypt(ciphertext) decrypt ciphertext with key=sha256(shared_key)
	sub decrypt {
		my iden $self=shift;
		unless (ref $self) { die "please use decrypt properly $!\n"; }
		my $ciphertext=shift || die "please specify a ciphertext to decrypt $!\n";

		my $cipher = Crypt::Rijndael->new(sha256($self->{shared_key}),Crypt::Rijndael::MODE_CTR());
		$cipher->set_iv($self->{_iv});
		my $plaintext  = $cipher->decrypt(_padd($ciphertext));
		return $plaintext;
	}
}
################################################################################################


#my $alice=iden->new("alice");
#my $bob=iden->new("bob");

#$alice->print_info;
#$bob->print_info;
#$alice->exchange($bob->{pk});
#$bob->exchange($alice->{pk});
#$alice->print_info;
#$bob->print_info;
#my $a=$alice->encrypt("hello world");
#print $bob->decrypt($a);

#send your keys
sub auth_exch0 {
	my $sock=shift||die "please pass the socket: $!\n";
	my $iden=shift||die "please pass the iden: $!\n";
	if(DEBUG) {
		print "\nSIGN={".$iden->sign($iden->{pk})."}KEY={".$iden->{pk}."};\n";
	}
	print $sock "SIGN={".$iden->sign($iden->{pk})."}KEY={".$iden->{pk}."};\r\n";
}

#receive the key and verify with yours
sub auth_exch1 {
	my $sock=shift||die "please pass the socket: $!\n";
	my $iden=shift||die "please pass the iden: $!\n";

	#Expecting auth information
	my $auth=<$sock>;
	if($auth=~m/SIGN=\{(.+)\}KEY=\{(.+)\};/) {
		if(DEBUG) {
			print "\nsignature=$1\nkey=$2\n";
		}
		if(!($iden->verify($1,$2))) {
			print $sock "ERROR={Signature agreement failed}";
			print "ERROR={Signature agreement failed}";
			close($sock);
			exit(-1);
		}
		$iden->exchange($2);
	} else {
		print $sock "ERROR={Unexpected authentication approach};GOT={$auth}";
		print "ERROR={Unexpected authentication approach};GOT={$auth}";
		close($sock);
		exit(-1);
	}

}

sub chat {
	my $sock=shift||die "please pass the socket: $!\n";
	my $iden=shift||die "please pass the iden: $!\n";
	print ": ";
	while(1) {
		my $i=<STDIN>;
		my $l=<$sock>;
		if(defined($l)) {
			my $len=length($l);
			$l=substr $l,0,$len-2;
			if(DEBUG) {print "read from socket :[".$l."]\n"; }
			if($l=~m/SIGN=\{(.+)\}MSG=\{(.*)\}KEY=\{(.+)\};/) {
				my $sig=$1;
				my $ciph=$2;
				my $key=$3;
				if(DEBUG) { print "sig=$sig\nciph=$ciph\nkey=$key\n"; }
				if(!($iden->verify($sig,$ciph.$key))) { #verify with old hmac_key
					print $sock "ERROR={Signature agreement failed}";
					print "ERROR={Signature agreement failed}";
					close($sock);
					exit(-1);
				}
				my $message=$iden->decrypt($ciph); #decrypt with old shared_key
				print "> $message";
				$iden->rekey; #rekey
				$iden->hmac_rekey; #rekey
				$iden->exchange($key); #add new public key

				my $msg="REKEY={".$iden->sign($iden->{pk})."}KEY={".$iden->{pk}."};\r\n"; #sign with new hmac_key
				if(DEBUG) {
					print $msg;
				}
				print $sock $msg;
				if(DEBUG) {$iden->print_info;}
			} elsif($l=~m/REKEY=\{(.+)\}KEY=\{(.+)\};/) {
				my $sig=$1;
				my $key=$2;
				if(DEBUG) { print "sig=$sig\nkey=$key\n"; }
				if(!($iden->verify($sig,$key))) { #verify with new hmac_key
					print $sock "ERROR={Signature agreement failed}";
					print "ERROR={Signature agreement failed}";
					close($sock);
					exit(-1);
				}
				$iden->exchange($key); #add new pubic key
				if(DEBUG) {$iden->print_info;}
			}
			print ": ";
		} elsif(defined($i)) {
			if($i=~/\\[exit|quit]/i) { close($sock); exit(0); }
			if(DEBUG) {print "read from stdin: $i\n"; }
			my $c=$iden->encrypt($i); #encrypt with the old shared_key
			$iden->rekey; #rekey
			my $msg="SIGN={".$iden->sign($c.($iden->{pk}))."}MSG={$c}KEY={".$iden->{pk}."};\r\n"; #sign with old hmac_key
			$iden->hmac_rekey; #rekey
			if(DEBUG) { print $msg; }
			print $sock $msg;
			if(DEBUG) {$iden->print_info;}
		} 
	}

}

sub server {
	my $port=shift||die "please speficy port for server $!\n";
	printf "Starting server on %s:%d\n",HOSTNAME,$port;


	my $sock = IO::Socket::INET->new( LocalHost=>HOSTNAME, LocalPort => $port, Proto => 'tcp', Listen=>1,ReuseAddr=>1)
			or die "Could not create socket: $!\n";
	
	my $client=$sock->accept();

	my $alice=iden->new("alice","killercodingninjas");
	if(DEBUG) {$alice->print_info;}

	print "Initial authentication...\n";
	auth_exch0($client,$alice); # send the signed key and public key
	auth_exch1($client,$alice); # receive the signed key and public key, and veify
	print "COMPLETE!\n";
	if(DEBUG) {$alice->print_info;}

	fcntl(STDIN,F_SETFL,O_NONBLOCK) or die "Could not set flags for STDIN: $!\n";
	fcntl($client,F_SETFL,O_NONBLOCK) or die "Could not set flags for $client: $!\n";
	sleep(1);
	chat($client,$alice);

	close($client);
}
sub client {
	my $host=shift||die "please specify hostname to client $!\n";
	my $port=shift||die "please specify port to client $!\n";
	print "Connecting to $host:$port...\n";

	my $sock = IO::Socket::INET->new( PeerAddr=>$host, PeerPort => $port, Proto => 'tcp') or die "Could not create socket: $!\n";

	my $bob=iden->new("bob","killercodingninjas");
	if(DEBUG) {$bob->print_info;}

	print "Initial authentication...\n";
	auth_exch1($sock,$bob);
	auth_exch0($sock,$bob);
	print "COMPLETE!\n";
	if(DEBUG) {$bob->print_info;}

	fcntl(STDIN,F_SETFL,O_NONBLOCK) or die "Could not set flags for STDIN: $!\n";
	fcntl($sock,F_SETFL,O_NONBLOCK) or die "Could not set flags for $sock: $!\n";
	sleep(1);
	chat($sock,$bob);


	close($sock);
}

sub help {
	print "$0 -[OPTIONS]\n"
	      ."OPTIONS\n"
	      ."server <port>      - start program in srver mode\n"
	      ."client <host:port> - start program connectin to server\n"
	      ."help		   - this\n";
}

sub invalidusage {
	print "Error: Invalid usage\n";
	help();
	exit(-1);
}

sub test {

	my $alice=iden->new("alice","killercodingninjas");
	my $bob=iden->new("bob","killercodingninjas");
	$alice->print_info; $bob->print_info;
	print "-"x30; print "\n";

	$bob->exchange($alice->{pk});
	$alice->exchange($bob->{pk});
	$alice->print_info; $bob->print_info;
	print "-"x30; print "\n";

	$alice->hmac_rekey;
	$bob->hmac_rekey;
	$alice->print_info; $bob->print_info;
	print "-"x30; print "\n";

	$alice->rekey;
	$bob->rekey;
	$alice->print_info; $bob->print_info;
	print "-"x30; print "\n";

	$bob->exchange($alice->{pk});
	$alice->exchange($bob->{pk});
	$alice->print_info; $bob->print_info;
	print "-"x30; print "\n";


}
	


if(!($#ARGV ge 0)) {
	invalidusage();
} else {
	if($ARGV[0] =~/server/i) {
		print "Starting server...";
		if($#ARGV ne 1) { invalidusage(); }
		server($ARGV[1]);
	} elsif($ARGV[0] =~/client/i) {
		print "Starting client...";
		if($#ARGV ne 1) { invalidusage(); }
		if($ARGV[1]=~/(.+):(\d+)/) {
			client($1,$2);
		} else { invalidusage(); }
	} elsif($ARGV[0]=~/help/i) {
		help();
		exit(0);
	} elsif($ARGV[0]=~/test/i) {
		test();
	}else { invalidusage(); }
}