#!/usr/bin/env ruby

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

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 to_s() @file.to_s end
	def basename() @file.basename end
	def dirname() @file.dirname end
	alias :to_path :file

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

class PublicKey
	Types = %w[sk-ecdsa-sha2-nistp256@openssh.com ecdsa-sha2-nistp256 ecdsa-sha2-nistp384 ecdsa-sha2-nistp521 sk-ssh-ed25519@openssh.com ssh-ed25519 ssh-dss ssh-rsa]
	RE = %r{
				(?<type> (?:#{Types.map{|t|Regexp.quote t}.join '|'}))
				\ +
				(?<pubkey> [a-zA-Z0-9/+]+ ) # base64 encoded
				\ +
				(?<comment> .*? ) # we use it as ident
				\ *
			}x
	Re = /\A#{RE}\z/

	attr :type, :key
	attr_accessor :comment

	def self.parse line
		m = Re.match line
		return nil  if m.nil?
		self.new m[:type], m[:pubkey], m[:comment]
	end

	def initialize type, key, comment
		die "Unknown PubKey type <#{type}>"  unless Types.include?( type)
		die "Not a (base64-encoded) PubKey <#{key}>"  if %r<[^a-zA-Z0-9/+]> =~ key
		@type, @key, @comment = type, key, comment
	end

	def to_s
		"#{@type} #{@key} #{@comment}"
	end

	def == other_pub
		other_pub.kind_of?( PublicKey) and @type == other_pub.type and @key == other_pub.key
	end
end

class AuthorizedKeys
	LINERE = %r{
				(?:(?<options>
					(?:[a-z-]+ (?:=\"[^"]*\")?)           # one option
				  (?:, (?:[a-z-]+ (?: = \"[^"]*\")?) )* # , more options
				)\ )?
				#{PublicKey::RE}
			}x
	LineRe = /\A#{LINERE}\z/
	attr_reader :file

	def self.parse line
		if m = LineRe.match( line)
			return PublicKey.new( m[:type], m[:pubkey], m[:comment]), m[:options]
		end
	end

	def initialize file
		@file = file
	end

	def each_line &block
		@file.each_line &block
	end

	def each_parsed_line &block
		return to_enum(__method__)  unless block_given?
		@file.each_line do |l|
			case l.chomp!
			when /\A\s*# *(.*)\z/
				yield :comment, l, $1
			when /\A\s*$/
				yield :empty, l
			else
				pk, o = self.class.parse( l)
				if pk
					yield :pub, l, o, pk
				else
					yield :unknown, l
				end
			end
		end
	end

	def each &block
		each_parsed_line {|t, *a| yield *a  if :pub == t }
	end

	def get ident
		each do |line, opts, pk|
			return opts, pk  if ident === pk.comment
		end
		nil
	end
	alias :[] :get

	def add! options, pub
		@file.open 'a+' do |f|
			f.puts "#{options} #{pub}"
		end
	end

	def delete! ident
		tmp = Pathname.new "#{@file}.tmp.#{Process.pid}"
		ret = 0
		tmp.open 'w' do |f|
			self.each_parsed_line do |t, l, *a|
				if :pub == t and a[1].comment == ident
					ret += 1
				else
					f.puts l
				end
			end
		end
		if 0 == ret
			tmp.unlink
		else
			@file.unlink
			tmp.rename @file
		end
		ret
	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, ident
		opts, entry = authkeysfile[ident]
		die "PubKey not found in authorized_keys: #{ident}"  unless entry
		if exist?
			entry == @file.read.chomp
		else
			@file.open( 'a+') {|f| f.puts entry }
			entry == PublicKey.parse( @file.read.chomp)
		end
	rescue Errno::ENOENT
		die "PubKey unusable"
	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:, pubsdir:
		#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 = pubsdir + '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

class Lock
	attr_reader :path, :file
	LOCK_T = [Fcntl::F_WRLCK, IO::SEEK_SET, 0, 0, 0].pack('s_s_q_q_q_')
	UNLOCK_T = [Fcntl::F_UNLCK, IO::SEEK_SET, 0, 0, 0].pack('s_s_q_q_q_')

	def initialize path
		@path = path
		@file = @path.open File::CREAT|File::RDWR, 0600
	end

	def exclusive!() die "Unable to get lock."  unless @file.fcntl Fcntl::F_SETLKW, LOCK_T end
	def release!() die "Unable to release lock."  unless @file.fcntl Fcntl::F_SETLKW, UNLOCK_T end

	def exclusive &block
		exclusive!
		yield
	ensure
		release!
	end
end

module REs
	HOSTNAME = /[a-zA-Z](?:[a-zA-Z0-9-]*[a-zA-Z0-9])?/
	Hostname = /\A#{HOSTNAME}\z/
	DOMAINNAME = /#{HOSTNAME}(?:\.#{HOSTNAME})*\.?/
	Domainname = /\A#{DOMAINNAME}\z/
	USERNAME = /[a-zA-Z0-9_]+/
	Username = /\A#{USERNAME}\z/
	USERIDENT = /#{USERNAME}@#{DOMAINNAME}/
	Userident = /\A#{USERIDENT}\z/

	# https://stackoverflow.com/a/17871737
	OCTET = /1?[0-9]?[0-9] | 2[0-4][0-9] | 25[0-5]/x
	IPV4ADDR = /#{OCTET}\.#{OCTET}\.#{OCTET}\.#{OCTET}/
	IPv4Addr = /\A#{IPV4ADDR}\z/

	IPV6SEG = /[0-9a-fA-Z]{1,4}/
	IPV6ADDR = %r<
		(?:#{IPV6SEG}:){7,7}#{IPV6SEG}|             # 1:2:3:4:5:6:7:8
		(?:#{IPV6SEG}:){1,7}:|                      # 1::                                 1:2:3:4:5:6:7::
		(?:#{IPV6SEG}:){1,6}:#{IPV6SEG}|            # 1::8               1:2:3:4:5:6::8   1:2:3:4:5:6::8
		(?:#{IPV6SEG}:){1,5}(?::#{IPV6SEG}){1,2}|   # 1::7:8             1:2:3:4:5::7:8   1:2:3:4:5::8
		(?:#{IPV6SEG}:){1,4}(?::#{IPV6SEG}){1,3}|   # 1::6:7:8           1:2:3:4::6:7:8   1:2:3:4::8
		(?:#{IPV6SEG}:){1,3}(?::#{IPV6SEG}){1,4}|   # 1::5:6:7:8         1:2:3::5:6:7:8   1:2:3::8
		(?:#{IPV6SEG}:){1,2}(?::#{IPV6SEG}){1,5}|   # 1::4:5:6:7:8       1:2::4:5:6:7:8   1:2::8
		#{IPV6SEG}:(?:(?::#{IPV6SEG}){1,6})|        # 1::3:4:5:6:7:8     1::3:4:5:6:7:8   1::8
		:(?:(?::#{IPV6SEG}){1,7}|:)|                # ::2:3:4:5:6:7:8    ::2:3:4:5:6:7:8  ::8       ::
		fe80:(?::#{IPV6SEG}){0,4}%[0-9a-zA-Z]{1,}|  # fe80::7:8%eth0     fe80::7:8%1  (link-local IPv6 addresses with zone index)
		::(?:ffff(?::0{1,4}){0,1}:){0,1}IPV4ADDR|   # ::255.255.255.255  ::ffff:255.255.255.255  ::ffff:0:255.255.255.255 (IPv4-mapped IPv6 addresses and IPv4-translated addresses)
		(?:#{IPV6SEG}:){1,4}:IPV4ADDR               # 2001:db8:3:4::192.0.2.33  64:ff9b::192.0.2.33 (IPv4-Embedded IPv6 Address)
	>x
	IPv6Addr = /\A#{IPV6ADDR}\z/

	IPADDR = /#{IPV6ADDR}|#{IPV4ADDR}/
	IPAddr = /\A#{IPADDR}\z/
end

class Main
	attr_reader :admin, :type, :ident, :principals, :info, :expire_in, :gracetime
	alias :admin? :admin

	def initialize *argv
		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]
		@admin, @type, @expire_in, @gracetime = false, type.to_sym, 12.weeks, 1.week
		prepare
		files
	end

	def prepare
		case @type
		when :host
			die "Invalid ident: #{@ident}"  unless REs::Domainname =~ @ident
			@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 REs::Userident =~ @ident
			@info ||= "user: #{@ident}"
			@principals = @principals.split ','

		when :admin
			die "Invalid ident: #{@ident}"  unless REs::Userident =~ @ident
			@info ||= "admin: #{@ident}"
			@type, @admin = :user, true
			@principals = @principals.split ','

		else
			die "Invalid type: #{@type}"
		end
	end

	attr_reader :lock, :pubsdir, :pubfile, :certfile, :cafile, :capubfile, :serialfile, :authkeysfile
	def files
		@config =
			begin
				JSON.parse Pathname.new( "~/.config/ssh-ca.conf").read
			rescue Errno::ENOENT
				{}
			rescue JSON::JSONError
				die "500, Internal server error"
			end
		@lock         = Lock.new Pathname.new( @config[:lockfile] || "/run/lock/ssh-ca.lock")
		@storage      = Pathname.new( @config[:storage] || "~/.local/sshca")
		@pubsdir     =
			if @config[:pubsdir] then Pathname.new @config[:pubsdir]
			else @storage + 'pubs'
			end.expand_path
		@pubfile      = PublicKeyFile.new @pubsdir + "#{@type}-#{@ident}.pub", @type, @ident
		@certfile     = CertifcateFile.new @pubsdir + "#{@type}-#{@ident}-cert.pub"
		@cafile       = @storage + "ca"
		@capubfile    = @storage + "ca.pub"
		@serialfile   = SerialFile.new @storage + "serial"
		@authkeysfile = AuthorizedKeys.new Pathname.new( @config[:authorized_keys_file] || "~/.ssh/authorized_keys").expand_path
		@binfile      = @config[:command] || __FILE__
	end

	def help msg=nil
		STDERR.puts msg  if msg
		conn = "#{ENV['LOGNAME']}@#{`hostname -f`}".chomp
		STDERR.puts <<-EOHELP.gsub( /^\s*#(.*)$/, admin? ? ":\\1" : "").sub(/\n+/, "\n").gsub( /^\s*:/, '')
			:Usage: ssh #{conn} renew [force] [show]
			:       ssh #{conn} show
			:       ssh #{conn} help
			:       ssh #{conn} jsonl
			#       ssh #{conn} keylist
			#       ssh #{conn} keyimport host|user|admin IDENT PRINCIPALS*
			#       ssh #{conn} keydelete IDENT
			:
			:[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.
			#keylist     Lists installed keys
			#keyimport host|user|admin IDENT PRINCIPALS*
			#            Import new key as host (a host-key) or user|admin (a user-key)
			#             Users and hosts are allowed to renew there keys.
			#             Admins are allowed to administrate keys (import/delete).
			#keydelete IDENT
			#            Deletes key of this IDENT
			:
			: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

	def show argv
		die "Unexpected addition parameters: #{argv.shelljoin}"  unless 1 == argv.length
		die "No certfile, yet: #{@certfile.basename}"  unless @certfile.exist?
		lock.exclusive do
			puts @certfile.read
		end
	end

	def renew argv
		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?

		lock.exclusive do
			@pubfile.prepare @authkeysfile, ident
			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, pubsdir: @pubsdir
			puts @certfile.read  if show
		end
	end

	def 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]

					lock.exclusive do
						@pubfile.prepare authkeyfile, ident
						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, pubsdir: @pubsdir,
							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
					end

				when :get
					lock.exclusive do
						STDOUT.puts [200, @certfile.metadata, @certfile.read].to_json
					end

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

	def keyimport argv
		help "Unknown command: #{argv[0]}"  unless admin
		type, ident, *principals = argv[1..-1].map(&:downcase)
		case type
		when 'host'
			die "Ident must be a domainname / hostname"  unless REs::Domainname =~ ident
			principals.each do |princ|
				case princ
				when REs::Domainname  then # ok
				when REs::IPv6Addr then # ok
				when REs::IPv4Addr then # ok
				else die "Principals must be domainnames or IP-addresses, not <#{princ}>"
				end
			end
		when 'user', 'admin'
			die "Ident must be user@hostname (lower cased)"  unless REs::Userident =~ ident
		else
			die "Type of key expected (host, user, admin) instead of <#{type}>."
		end

		str = STDIN.readline.chomp
		pk, _ = PublicKey.parse( str)
		die "Cannot parse public key <#{str}>."  unless pk
		pk.comment = ident
		lock.exclusive do
			@authkeysfile.add! "restrict,command=\"#{__FILE__} #{type} #{ident} #{principals.join ','}\"", pk
		end
	end

	def keydelete argv
		help "Unknown command: #{argv[0]}"  unless admin
		ident = argv[1]
		lock.exclusive do
			@authkeysfile.delete! ident
		end
	end

	def keylist argv
		help "Unknown command: #{argv[0]}"  unless admin
		lock.exclusive do
			@authkeysfile.each_line {|l| STDOUT.print l }
		end
	end

	def exec cmd
		help  unless cmd
		die "Unallowed char in arguments [\\0]"  if cmd.include? 0.chr
		argv = cmd.shellsplit
		case argv[0]
		when *%w[h help] then help
		when *%w[s show] then show argv
		when *%w[r renew] then renew argv
		when 'jsonl' then jsonl argv
		when 'keyimport' then keyimport argv
		when 'keydelete' then keydelete argv
		when 'keylist' then keylist argv
		else help "Unknown command: #{argv[0]}"
		end
	end
end

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

	main = Main.new *ARGV
	main.exec ENV['SSH_ORIGINAL_COMMAND']

rescue SystemExit
	raise

rescue Die
	STDERR.puts $!
	exit 1

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