#!/usr/bin/perl

# Usage: generate_keys_and_flag.pl output-directory [number-of-keys]
#
# output-directory must exist and not contain any files, public keys
# and flag will be stored there. This is for revealing to the contester.
# 
# number-of-keys defaults to 12. In any case, only one pair will be
# vulnerable to gcd attack

use strict;
use Math::BigInt;
use Convert::ASN1;
use MIME::Base64;
use File::Temp;

# Default configuration
my $keys = 12;

die "Usage: $0 output-dir [iterations]" if $#ARGV lt 0;

# check that destination exists and is empty
opendir(my $dh, $ARGV[0]) or die "$ARGV[0]: $!";
die "$ARGV[0] not empty" if scalar(grep { $_ ne "." && $_ ne ".." } readdir($dh)) > 0;

my $flagfile = "$ARGV[0]/encrypted_flag";
$keys = $ARGV[1] if $#ARGV > 0;
die "Need at least two keys for the task" if $keys < 2;

# ASN.1 template of OpenSSL private key file
my $asn = Convert::ASN1->new;
$asn->prepare(q<
    SEQUENCE {
      id INTEGER,
      n INTEGER,
      e INTEGER,
      d INTEGER,
      p INTEGER,
      q INTEGER,
      e1 INTEGER,
      e2 INTEGER,
      coef INTEGER
    }
  >);

my @allkeys = ();

print "Generating $keys RSA keypairs\n";

# generate good keys and save to files with random names
for (1..$keys) {
	my $fh = File::Temp->new(DIR => $ARGV[0], UNLINK => 0) or die $ARGV[0].": $!";
	system("openssl genrsa -f4 -out ".$fh->filename." 2>/dev/null") == 0 or die "openssl: $!";
	push @allkeys, $fh->filename;
	close($fh);
}

# pick randomly two candidates for weak keys from last 10% of the file list
my @sorted = reverse sort @allkeys;
my $key1 = int(rand($keys/20));
my $key2 = $key1 + int(rand($keys/20));
$key2 = 0 if $key2 >= $keys;
$key2 = 1 if $key2 == $key1;

# read from the chosen key and return Convert::ASN1 object
sub readkey {
	open(my $fh, "<", $_[0]) or die "$_[0]: $!";
	my $der = do { local $/; decode_base64(<$fh> =~ s/^-.*?\n//gmr); };
	close($fh);
	my $key = $asn->decode($der);
	die "$$_[0]: ".$asn->error() if not $key;
	return $key;
}

# replace one prime in the second weak key
my $p=readkey($sorted[$key1])->{'p'};
my $q=readkey($sorted[$key2])->{'p'};

# compute new key
my $p1=$p->copy()->bsub(1);
my $q1=$q->copy()->bsub(1);

# e is fixed in all keys generated by openssl
my $e=Math::BigInt->new(65537);

die "e and lambda(n) are not coprimes, pick another keys" if (
			$e->copy()->bgcd($p1->copy()->blcm($q1)) != 1);

my $n=$p->copy()->bmul($q);
my $d=$e->copy()->bmodinv($p1->copy()->bmul($q1));
my $e1=$d->copy()->bmod($p1);
my $e2=$d->copy()->bmod($q1);
my $qinv = $q->copy()->bmodinv($p);

print "Creating weak keys $sorted[$key1] and $sorted[$key2]\n";

# overwrite one key
open(my $fh, ">", $sorted[$key2]) or die "$sorted[$key2]: $!";
print $fh "-----BEGIN RSA PRIVATE KEY-----\n";
print $fh encode_base64($asn->encode(id => 0, n=>$n, e=>$e, d=>$d, p=>$p, q=>$q, e1=>$e1, e2=>$e2, coef=>$qinv));
print $fh "-----END RSA PRIVATE KEY-----\n";
close($fh);

# create flag
open(my $ph, "-|", "ps ax|openssl dgst -hex -sha1 -r") or die "$!";
my $flag = "The flag is: ECSC{" . do { local $/; <$ph> =~ s/ \*stdin.*//sr; } . "}\n";
close($ph);

open($ph, "|-", "openssl rsautl -encrypt -inkey $sorted[$key1] > $flagfile") or die "openssl: $!";
print $ph $flag;
close($ph) or die "openssl: $!";
print "Flag saved to $flagfile\n";

# extract public keys and remove private keys
print "Removing private keys\n";
system("openssl rsa -in $_ -pubout -out $_.pub 2>/dev/null && rm $_") == 0 or die "openssl: $!" for (@allkeys);

print "$flag\n";