#!/usr/bin/env ruby

require 'time'
require 'pathname'
require 'shellwords'
require 'getoptlong'
require 'active_support/time'
require 'json'

def help msg=nil
	STDERR.puts msg  if msg
	STDERR.puts <<-EOHELP.gsub( /^\s*:/, '')
		:Usage: ssh .... renew [force] [show]
		:       ssh .... show
		:       ssh .... help
		:       ssh .... jsonl
		:
		:[h]elp
		:[r]enew     Renews (if expires in graceperiod) or creates (if not existing) certificate
		: [f]orce    forces renewal, also if valid for graceperiod
		:[s]how      Shows current certificate
		:jsonl       Expect commands via STDIN as JSONL and responses via STDOUT in JSONL
		:             Request format: [0, CMD, ARGS...]
		:             Response format: [SCODE, BODY]
		:             SCODEs are like HTTP.
		:
		:possible CMD:
		:<renew[, {<force: true|false>}]> => <RENEWED, DATA, CONTENT>
		:<get>                            => <DATA, CONTENT>
		:
		:  RENEWED indicates if file was renewed (true|false).
		:  CONTENT is the content of the certfile (String).
		:  DATA is the parsed file (ssh-keygen -L) (Hash).
		EOHELP
	exit 1
end

class Die <Exception
end

module Kernel
	def popen *a, **o, &e
		a = a.flatten.compact.map &:to_s
		IO.popen a, **o, &e
	end

	def die str
		raise Die, str
	end
end

module FileHelper
	attr_reader :file

	def exist?
		@file.file?
	end

	def read
		@file.read
	end

	def open *a, **o, &e
		if block_given?
			@file.open *a, **o, &e
		else
			@file.open *a, **o
		end
	end

	def to_s
		@file.to_s
	end
end

class PublicKeyFile
	include FileHelper
	attr_reader :type, :ident

	def initialize file, type, ident
		@file, @type, @ident = file, type, ident
	end

	def prepare authkeysfile
		return true  if exist?
		authkeysfile.each_line do |l|
			case l.chomp
			when /\A\s*#/ # ignore - comment
			when /\A\s*$/ # ignore - empty line
			when /\Assh-/ # ignore - blank keys without command
			when / (ssh-(?:rsa|ed25519) +[^ ]+ +#{Regexp.quote @ident}) *\z/
				pub = $1
				@file.open( 'a+') {|f| f.puts pub }
			end
		end
		exist?
	end
end

class CertifcateFile
	include FileHelper
	attr_reader :metadata

	def initialize file
		@file = file
	end

	def parse
		r = {}
		popen %w[ssh-keygen -L -f], @file do |f|
			mode = nil
			f.each_line do |l|
				case l.chomp
				when /^        Valid: from ([^ ][^ ]*) to ([^ ][^ ]*)$/
					r[:valid] = Time.parse( $1) ... Time.parse( $2)
				when /^        Type: ([^ ]+) ([^ ]+) certificate$/
					r[:type]  = [$1, $2]
				when /^        Public key: ([^ ]+) (.*)$/
					r[:pubkey]= [$1, $2]
				when /^        Signing CA: ([^ ]+) (.*)$/
					r[:signca]= [$1, $2]
				when /^        Key ID: (.*)$/
					r[:keyid] = $1.sub /\A"?(.*?)"?\z/, '\1'
				when /^        Serial: (\d+)$/
					r[:serial]= $1.to_i
				when /^        Principals: \(node\)$/
					r[:principals] = []
				when /^        Principals:\s*$/
					mode = :principals
					r[:principals] = []
				when /^        Extensions: \(none\)$/
					r[:extensions] = []
				when /^        Extensions:\s*$/
					mode = :extensions
					r[:extensions] = []
				when /^        Critical Options: \(none\)$/
					r[:critopts] = []
				when /^        Critical Options:\s*$/
					mode = :critopts
					r[:critopts] = []
				when /^                ([^ ].*)$/
					r[mode].push $1
				else
				end
			end
		end
		@metadata = r
	end

	def [] key
		(@metadata || parse)[key]
	end
	alias to_hash []

	def valid_at? ts
		exist? and self[:valid].include?( ts)
	end

	def create pubfile, type:, valid:, serial:, cafile:, info:, principals:, sshcadir:
		#sign_user_pub opts
		err =
			popen( %w[ssh-keygen -V], "#{valid.begin.strftime '%Y%m%d'}:#{valid.end.strftime '%Y%m%d'}",
						(:host == type ? '-h' : nil), '-s', cafile,
						'-P', '', '-z', serial, '-I', info, '-n', principals.join(','), pubfile,
					err: %i[child out]
				) {|l| l.read }
		die "Creating certificate failed [#{$?.inspect}]: #{err}"  unless 0 == $?.exitstatus
		serialkeyfile = sshcadir + 'serials' + serial.to_s
		serialkeyfile.open( 'a+') {|f| f.write self.read }
		self
	end
end

class SerialFile
	include FileHelper

	def initialize file
		@file = file
	end

	def read
		open 'a+' do |fh|
			fh.read.to_i
		end
	end
	alias to_i read
	alias get read

	def increment!
		open 'a+' do |fh|
			(1+fh.read.to_i).tap do |val|
				fh.rewind
				fh.truncate 0
				fh.print val
			end
		end
	end
	alias inc! increment!
end


#################################################
begin

die "2..4 arguments (type ident [principals [info]]) expected.  Provided count of arguments: #{ARGV.length}"  unless (2..4) === ARGV.length
type, ident, principals, info = ARGV[0..3]
expire_in, gracetime = 12.weeks, 1.week
type = type.to_sym
case type
when :host
	die "Invalid ident: #{ident}"  unless ident =~ /\A[a-z][a-z_0-9]*\z/i
	expire_in, gracetime = 360.days, 30.days
	info ||= "host: #{ident}"
	principals =
		([ident] +
			principals.split( ',').map do |e|
				if /:/ =~ e or /\A[0-9.]+\z/ =~ e or /\A[^.]+\z/ =~ e
					e
				else
					e = e.gsub /\.?\z/, ''
					[e, "#{e}."]
				end
			end
		).flatten.compact.map do |e|
			if /\A\[/ =~ e
				e
			else
				[e, "[#{e}]:19"]
			end
		end.flatten.compact.uniq
when :user
	die "Invalid ident: #{ident}"  unless ident =~ /\A[a-z][a-z_0-9]*@[a-z][a-z_0-9]*\z/i
	info ||= "user: #{ident}"
	principals = principals.split ','
else
	die "Invalid type: #{type}"
end

sshcadir     = Pathname.new( "~/.ssh-ca").expand_path
pubfile      = PublicKeyFile.new sshcadir + "#{type}-#{ident}.pub", type, ident
certfile     = CertifcateFile.new sshcadir + "#{type}-#{ident}-cert.pub"
cafile       = sshcadir + "ca"
capubfile    = sshcadir + "ca.pub"
serialfile   = SerialFile.new sshcadir + "serial"
authkeysfile = Pathname.new( "~/.ssh/authorized_keys").expand_path

help  unless ENV['SSH_ORIGINAL_COMMAND']
die "Unallowed char in arguments [\\0]"  if ENV['SSH_ORIGINAL_COMMAND'].include? 0.chr
ARGV.replace ENV['SSH_ORIGINAL_COMMAND'].shellsplit
case ARGV[0]
when *%w[h help]
	help

when *%w[s show]
	die "Unexpected addition parameters: #{ARGV.shelljoin}"  unless 1 == ARGV.length
	die "No certfile, yet: #{certfile.basename}"  unless certfile.exist?
	puts certfile.read

when *%w[r renew]
	begin
		unexpected = ARGV[1..-1] - %w[s show f force]
		die "Unknown argument: #{unexpected.shelljoin}"  unless unexpected.empty?
	end
	show  = ! (ARGV[1..-1] & %w[s show]).empty?
	force = ! (ARGV[1..-1] & %w[f force]).empty?

	pubfile.prepare authkeysfile
	die "Pubfile or pub in authorized_keys for #{ident} not found."  unless pubfile.exist?

	if !force && certfile.exist? && certfile.valid_at?( gracetime.from_now) && certfile[:principals].sort == principals.sort
		STDERR.print "Certificate [#{certfile[:serial]}|#{certfile[:keyid]}] valid between #{certfile[:valid]} as #{certfile[:principals].sort.join ', '}\n"
		STDERR.flush
		STDOUT.print "#{certfile.read.chomp}\n"  if show
		exit 0
	end

	STDERR.puts "Renew certificate for #{ident} as #{principals.sort.join ', '}"
	certfile.create pubfile, type: type, valid: -5.minutes.ago..expire_in.from_now,
		serial: serialfile.increment!, cafile: cafile, info: info, principals: principals, sshcadir: sshcadir
	puts certfile.read  if show

when 'jsonl'
	STDIN.each_line do |line|
		begin
			line = JSON.parse line, symbolize_names: true
			cmd, args  = line[0].to_sym, line[1..-1]
			case cmd
			when :renew
				die "renew expects no or a Hash as args."  unless (0..1).include?( args.length) && args.kind_of?( Hash)
				args ||= {}
				args[:force] = ! ! args[:force]

				pubfile.prepare
				die "Pubfile or pub in authorized_keys for #{ident} not found."  unless pubfile.exist?

				if !args[:force] && certfile.valid_at?( gracetime.from_now)
					STDOUT.puts [200, false, certfile.metadata, certfile.read].to_json
				end

				certfile.create pubfile, type: type, sshcadir: sshcadir,
					valid: -5.minutes.ago..expire_in.from_now, serial: serialfile.increment!,
					cafile: cafile, info: info, principals: principals
				STDOUT.puts [200, true, certfile.metadata, certfile.read].to_json

			when :get
				STDOUT.puts [200, certfile.metadata, certfile.read].to_json

			else
				die "Unknown command."
			end
		rescue Object
			STDOUT.puts [500, $!.class.name, $!.to_s].to_json
		end
	end

else
	help "Unknown command: #{ARGV[0]}"
end

rescue SystemExit
	raise
rescue Die
	STDERR.puts $!
	exit 1
rescue Object
	STDERR.puts "#$! (#{$!.class})", *$!.backtrace
	exit 2
end