#!/usr/bin/env ruby
# dhkex -- Diffie-Hellman key exchange
require 'base64'
require 'openssl'
require 'optparse'

def raw_output(raw)
    return Base64.strict_encode64(raw)
end
def bn_output(bn)
    return nil if !bn
    raw = [bn.to_s(16)].pack("H*")
    return raw_output(raw)
end

def raw_input(str)
    return nil if !str
    str = File.read($1) if str =~ /^@(.+)/
    str.strip!
    return nil if str == "" || str == "nil"
    return Base64.decode64(str)
end
def bn_input(str)
    raw = raw_input(str)
    return nil if !raw
    hex = raw.unpack("H*")[0]
    return OpenSSL::BN.new(hex, 16)
end

def buf_to_hex(buf)
    buf.unpack("H*")[0]
end

def color(str, fmt)
    if STDOUT.tty?
        "\e[#{fmt}m#{str}\e[m"
    else
        str
    end
end

def ask(str)
    if STDIN.tty?
        STDOUT.print "#{str} "
        STDOUT.flush
    end
    return STDIN.gets
end

curve_aliases = {
    "P-256" => "prime256v1",
    "P-384" => "secp384r1",
}
curve_name = nil
default_curve = "prime256v1"
my_private_str = nil
other_public_str = nil
show_private = false
ask_my_private = false
ask_other_public = true
secret_file = nil
use_hkdf = true

OptionParser.new do |opts|
    opts.banner = "Usage: dhkex [options]"
    opts.separator("")
    opts.on("-y", "--their-pubkey KEY", String, "Other party's public key") do |s|
        other_public_str = s
        ask_other_public = false
    end
    opts.on("-o", "--write-secret FILE", String, "Write shared secret to file") do |s|
        secret_file = s
    end
    opts.separator("")
    opts.separator("Advanced options:")
    opts.separator("")
    opts.on("-b", "--curve NAME", String, "EC curve name (default: #{default_curve})") do |s|
        curve_name = curve_aliases[s] || s
    end
    opts.on("-R", "--resumable", "Display our private key") do |t|
        show_private = t
    end
    opts.on("-r", "--resume", "Ask for private key") do |t|
        ask_my_private = true
    end
    opts.separator("")
    opts.on("-g", "--generate-privkey", "Display private key and exit") do
        show_private = true
        ask_other_public = false
    end
    opts.on("-k", "--resume-with-privkey KEY", String, "Reuse old private key") do |s|
        my_private_str = s
        ask_my_private = false
    end
end.parse!

if !show_private and !other_public_str and !ask_other_public
    abort "error: specifying '-y nil' without '-R' does nothing useful"
end
#if (my_private_str or ask_my_private) and curve_name
#    abort "error: loading existing private key conflicts with specifying curve"
#end

# load our private key

if ask_my_private
    my_private_str = ask("my private key?") or exit
end
if my_private_str
    my_pkey = OpenSSL::PKey.read(raw_input(my_private_str))
    if curve_name
        curve = OpenSSL::PKey::EC::Group.new(curve_name)
        if curve != my_pkey.group
            abort "error: specified private key does not match '-b #{curve_name}'"
        end
    end
else
    my_pkey = OpenSSL::PKey::EC.generate(curve_name || default_curve)
end
if show_private
    puts "my private key: " + color(raw_output(my_pkey.to_der), "31")
end
puts "my public key: " + color(bn_output(my_pkey.public_key.to_bn), "32")

# load peer public key

if !other_public_str
    if !ask_other_public
        exit
    end
    other_public_str = ask("their public key?") or exit
end
other_public_bn = bn_input(other_public_str) or exit
other_public_point = OpenSSL::PKey::EC::Point.new(my_pkey.group, other_public_bn)

# compute shared secret

ikm = my_pkey.dh_compute_key(other_public_point)
nbits = ikm.length * 8

if use_hkdf
    secret = OpenSSL::KDF.hkdf(ikm, salt: "",
                                    info: "",
                                    # all the usual curves "offer n/2 bits of
                                    # security" so this seems appropriate
                                    length: ikm.length / 2,
                                    hash: OpenSSL::Digest::SHA256.new)
    nbits = secret.length * 8
    puts "\e[2musing HKDF-SHA256 to derive shared secret from DH result\e[m"
else
    secret = ikm
    puts "\e[1mnote: DH result should never be used raw, only as input to a hash or HKDF\e[m"
end

if secret_file
    File.open(secret_file, "wb") do |f|
        f.write(secret)
    end
    puts "shared secret (#{nbits} bits) written to \"#{secret_file}\""
else
    puts "shared secret (#{nbits} bits): " + color(Base64.strict_encode64(secret), "36")
end

# vim: ts=4:sw=4:et
