#!/usr/bin/env ruby
require 'rubygems'
require 'r509'
require 'r509/trollop'
require 'io/console'

opts = R509::Trollop.options do
  opt :interactive, "Interactive CSR/self-signed certificate generation. Overrides all flags other than keyout and out."
  opt :subject, "X509 subject / delimited. Example: /CN=test.com/O=Org/C=US/ST=Illinois/L=Chicago", :type => :string
  opt :san, "Subject Alternative Name Example: test.com,*.test.com", :type => :string
  opt :message_digest, "Message digest to use. sha1, sha224, sha256, sha384, sha512, md5", :type => :string, :default => 'sha256'
  opt :duration, "Self-sign the certificate with the duration (in days) specified.", :type => :integer
  opt :bits, "Bit length of generated key. Ignored for EC.", :type => :integer, :default => 2048
  opt :curve_name, "Name of elliptic curve to use. Only used for EC.", :type => :string, :default => 'secp384r1'
  opt :keyout, "File name to save generated key.", :type => :string
  opt :out, "File name to save generated CSR or self-signed certificate", :type => :string
  opt :type, "Type of key to generate. RSA/DSA/EC", :type => :string, :default => "RSA"
  opt :password, "Password to encrypt generated key", :type => :string
  if RUBY_PLATFORM.match('darwin')
    opt :clipboard, "Copy CSR or certificate to the clipboard", :default => false, :short => :p
  end
  version "r509 #{R509::VERSION}"
end

opts[:duration] = opts[:duration].to_i
subject = []

if opts[:interactive] == true || opts[:subject].nil?
  if opts[:type].upcase == "RSA" || opts[:type].upcase == "DSA"
    print "CSR Bit Length (2048):"
    bit_length = gets.chomp
    opts[:bits] = (bit_length.to_i > 0) ? bit_length.to_i : 2048
  elsif opts[:type].upcase == "EC"
    print "Curve Name (secp384r1):"
    curve_name = gets.chomp
    opts[:curve_name] = (!curve_name.empty?) ? curve_name : 'secp384r1'
  else
    puts "Invalid key type specified. RSA/DSA/EC"
    exit
  end

  print "Message Digest (#{R509::MessageDigest::DEFAULT_MD}):"
  md = gets.chomp
  opts[:message_digest] = case md
    when 'sha1' then 'sha1'
    when 'sha224' then 'sha224'
    when 'sha256' then 'sha256'
    when 'sha384' then 'sha384'
    when 'sha512' then 'sha512'
    when 'md5' then 'md5'
    else R509::MessageDigest::DEFAULT_MD
  end

  print "C (US): "
  c = gets.chomp
  c = c.empty? ? 'US' : c
  subject.push ['C', c]

  print "ST (Illinois): "
  st = gets.chomp
  st = st.empty? ? 'Illinois' : st
  subject.push ['ST', st]

  print "L (Chicago): "
  l = gets.chomp
  l = l.empty? ? 'Chicago' : l
  subject.push ['L', l]

  print "O (r509 LLC): "
  o = gets.chomp
  o = o.empty? ? 'r509 LLC' : o
  subject.push ['O', o]

  print "OU (null by default): "
  ou = gets.chomp
  unless ou.empty?
    subject.push ['OU', ou]
  end

  print "CN: "
  subject.push ['CN', gets.chomp]
  print "SAN Domains (comma separated):"
  opts[:san] = gets.chomp

  print "Self-signed cert duration in days (null disables self-sign):"
  opts[:duration] = gets.chomp.to_i

  print "Password to encrypt generated key (empty disables encryption):"
  password = STDIN.noecho(&:gets).chomp
  puts ""
  unless password.empty?
    print "Retype password:"
    password_confirm = STDIN.noecho(&:gets).chomp
    puts ""
    unless password == password_confirm
      puts "Passwords do not match."
      exit
    end
    opts[:password] = password
  end
else
  opts[:subject].chomp.split('/').each do |item|
    if item != ''
      subject.push item.split('=')[0..1]
    end
  end
end

csr_or_cert = csr = R509::CSR.new(
  :subject => subject,
  :bit_length => opts[:bits],
  :type => opts[:type].upcase,
  :curve_name => opts[:curve_name],
  :san_names => (opts[:san] || "").split(',').map { |domain| domain.strip },
  :message_digest => opts[:message_digest]
)

# for self signed, outputting the cert (not the csr)
selfsign = opts[:duration]
if selfsign > 0
  csr_or_cert = R509::CertificateAuthority::Signer.selfsign(
    :csr => csr,
    :not_after => Time.now.to_i + 86400 * selfsign,
    :message_digest => opts[:message_digest]
  )
end

if opts[:keyout]
  if opts[:password]
    csr.key.write_encrypted_pem(opts[:keyout], "aes256", opts[:password])
  else
    csr.key.write_pem(opts[:keyout])
  end
else
  if opts[:password]
    puts csr.key.to_encrypted_pem("aes256", opts[:password])
  else
    puts csr.key.to_pem
  end
end

if opts[:out]
  csr_or_cert.write_pem(opts[:out])
else
  puts csr_or_cert.to_pem
end

puts csr_or_cert.subject
puts "SAN(s): #{csr_or_cert.san.names.map { |n| n.value }.join(", ")}" if csr_or_cert.san

if opts[:clipboard]
  IO.popen('pbcopy', 'w').puts csr_or_cert
end