#!/usr/bin/env perl
# kc -- Kerberos credential cache juggler
# For cache switching to work, kc.sh must be sourced.
use warnings;
use strict;
use feature qw(say state);
use English;
use File::Basename;
use File::stat;
use File::Temp qw(tempfile);
use Nullroute::Lib;

my $rundir;
my $ccprefix;
my $runprefix;
my $cccurrent;
my $ccenviron;
my $ccdefault;
my $cccdir;
my $cccprimary;
my @caches;

my $can_switch = 1;

my %COLORS = (
	reset		=> "0",

	empty		=> "35",
	normal		=> "0",
	expiring	=> "33",
	expired		=> "31",

	empty_active	=> "95",
	normal_active	=> "32",
	expiring_active	=> "33",
	expired_active	=> "91",
);

sub usage {
	say for
	"Usage: $::arg0",
	"       $::arg0 <name>|\"@\" [kinit_args]",
	"       $::arg0 <number>",
	"       $::arg0 {list|slist}",
	"       $::arg0 purge",
	"       $::arg0 destroy <name|number>...";
}

sub _debugvar {
	my ($var, $val) = @_;
	@_ = ($var."='".($val//"")."'");
	goto &_debug;
}

sub fmt {
	my ($color, $str, $fmt) = @_;
	if (!$str) {
		$str = "";
	}
	if ($fmt) {
		$str = sprintf($fmt, $str);
	}
	if ($str && $color) {
		return $color.$str.$COLORS{reset};
	}
	return $str;
}

sub which {
	my ($name) = @_;
	state %paths;

	return $name if $name =~ m|/|;

	if (!exists $paths{$name}) {
		($paths{$name}) = grep {-x}
				map {"$_/$name"}
				map {$_ || "."}
				split(/:/, $ENV{PATH});
	}

	return $paths{$name};
}

sub run_proc {
	my (@argv) = @_;

	$argv[0] = which($argv[0]);

	return system(@argv);
}

sub read_proc {
	my (@argv) = @_;
	my ($proc, $pid, $output);

	$argv[0] = which($argv[0]);

	$pid = open($proc, "-|") // _die("could not fork: $!");

	if (!$pid) {
		no warnings;
		open(STDERR, ">/dev/null");
		exec({$argv[0]} @argv);
		_die("could not exec $argv[0]: $!");
	}

	chomp($output = <$proc> // "");
	close($proc);

	return $output;
}

sub read_file {
	my ($path) = @_;
	my $output;

	open(my $file, "<", $path) or _die("could not open $path: $!");
	chomp($output = <$file>);
	close($file);

	return $output;
}

sub get_keytab_path {
	my ($princ) = @_;

	# TODO: depend on hostname
	"$ENV{HOME}/Private/keys/login/krb-$princ.keytab";
}

sub read_aliases_from_file {
	my ($path, $aliases) = @_;

	_debug("reading aliases from '$path'");
	open(my $file, "<", $path) or return;
	while (my $line = <$file>) {
		next if $line =~ /^#/;
		chomp $line;
		my ($alias, @args) = split(/\s+/, $line);
		if (@args) {
			my %vars = (
				PRINCIPAL => $args[0],
				KEYTAB => get_keytab_path($args[0]),
			);
			for (@args) {
				s|^~/|$ENV{HOME}/|;
				s|\$\{([A-Z]+)\}|$vars{$1}|g;
			}
			$aliases->{$alias} = \@args;
		} else {
			warn "$path:$.: not enough parameters\n";
		}
	}
	close($file);
}

sub read_aliases {
	my @paths = (
		"$ENV{HOME}/.config/nullroute.eu.org/k5aliases",
		"$ENV{HOME}/.dotfiles/k5aliases",
		"$ENV{HOME}/lib/dotfiles/k5aliases",
		"$ENV{HOME}/lib/k5aliases",
	);
	my %aliases;

	for my $path (@paths) {
		read_aliases_from_file($path, \%aliases);
	}

	return %aliases;
}

sub enum_ccaches {
	my @ccaches;

	open(my $proc, "-|", which("pklist"), "-l", "-N")
	or _die("could not run 'pklist': $!");
	push @ccaches, grep {chomp or 1} <$proc>;
	close($proc);

	# traditional

	push @ccaches,	map {"FILE:$_"}
			grep {
				my $st = stat($_);
				-f $_ && $st->uid == $UID
			}
			glob("/tmp/krb5cc*");

	# new

	if (-d "$rundir/krb5cc") {
		push @ccaches,	map {"DIR::$_"}
				glob("$rundir/krb5cc/tkt*");

		push @ccaches,	map {"DIR::$_"}
				glob("$rundir/krb5cc_*/tkt*");
	}

	# Heimdal kcmd

	if (-S "/var/run/.heim_org.h5l.kcm-socket") {
		push @ccaches, "KCM:$UID";
	}

	# kernel keyrings

	my @keys = uniq map {split} grep {chomp or 1}
		   qx(keyctl rlist \@s 2>/dev/null),
		   qx(keyctl rlist \@u 2>/dev/null);
	for my $key (@keys) {
		my $desc = read_proc("keyctl", "rdescribe", $key);
		if ($desc =~ /^keyring;.*?;.*?;.*?;(krb5cc\.*)$/) {
			push @ccaches, "KEYRING:$1";
		}
	}

	# filter out invalid ccaches

	@ccaches = grep {run_proc("pklist", "-q", "-c", $_) == 0} @ccaches;

	# special ccaches (never filtered)

	my $have_current = grep {$_ eq $cccurrent} @ccaches;
	if (!$have_current) {
		push @ccaches, $cccurrent;
	}

	if (length $ccenviron) {
		my $have_environ = grep {ccache_is_environ($_)} @ccaches;
		if (!$have_environ) {
			push @ccaches, $ccenviron;
		}
	} else {
		my $have_default = grep {ccache_is_default($_)} @ccaches;
		if (!$have_default) {
			push @ccaches, $ccdefault;
		}
	}

	@ccaches = uniq sort @ccaches;

	return @ccaches;
}

sub expand_ccname {
	my ($name) = @_;
	for ($name) {
		if (m|^new$|) {
			my (undef, $path) = tempfile($ccprefix."XXXXXX", OPEN => 0);
			return "FILE:$path";
		}
		elsif (m|^@?$|) {
			return $ccdefault;
		}
		elsif (m|^KCM$|i) {
			return "KCM:$UID";
		}
		elsif (m|^\d{1,2}$|) {
			my $i = int $_;
			if ($i > 0 && $i <= @caches) {
				return $caches[$i - 1];
			}
		}
		elsif (m/^\^\^(.+)$/) {
			return "KEYRING:$1";
		}
		elsif (m/^\^(.+)$/) {
			return "KEYRING:persistent:$UID:$1";
		}
		elsif (m/^\^$/) {
			return "KEYRING:persistent:$UID";
		}
		# +foo
		elsif (m|^\+$|) {
			return "DIR:$cccdir";
		}
		elsif (m|^\+(.*)$|) {
			return "DIR::$cccdir/tkt$1";
		}
		# :foo/bar
		elsif (m|^:(.+)/$|) {
			return "DIR:$runprefix"."_$1";
		}
		elsif (m|^:(.+)/(.+)$|) {
			return "DIR::$runprefix"."_$1/tkt$2";
		}
		# :foo
		elsif (m|^:$|) {
			return "DIR:$runprefix";
		}
		elsif (m|^:(.+)$|) {
			return "DIR::$runprefix/tkt$1";
		}
		# any
		elsif (m|:|) {
			return $_;
		}
		elsif (m|/|) {
			return "FILE:$_";
		}
		else {
			return "FILE:$ccprefix$_";
		}
	}
}

sub collapse_ccname {
	my ($name) = @_;
	for ($name) {
		if ($_ eq $ccdefault) {
			return "@";
		}
		elsif (m|^DIR::\Q$runprefix\E_(.+)/tkt(.*)$|) {
			return ":$1/$2";
		}
		elsif (m|^DIR::\Q$runprefix\E/tkt(.*)$|) {
			return ":$1";
		}
		elsif (m|^DIR::\Q$cccdir\E/tkt(.*)$|) {
			return "+$1";
		}
		elsif (m|^FILE:\Q$ccprefix\E(.*)$|) {
			return $1;
		}
		elsif (m|^FILE:(/.*)$|) {
			return $1;
		}
		#elsif ($_ eq "API:$principal") {
		#	return "API:";
		#}
		elsif ($_ eq "KCM:$UID") {
			return "KCM";
		}
		elsif ($_ eq "KEYRING:persistent:$UID") {
			return "^";
		}
		elsif (m|^KEYRING:persistent:$UID:(.+)$|) {
			return "^$1";
		}
		elsif (m|^KEYRING:(.*)$|) {
			return "^^$1";
		}
		else {
			return $_;
		}
	}
}

sub cmp_ccnames {
	my ($a, $b) = @_;
	my $primary = "tkt";

	$a = "FILE:$a" unless $a =~ /:/;
	$b = "FILE:$b" unless $b =~ /:/;

	return 1 if $a eq $b;

	if ($a =~ /^DIR:([^:].*)$/) {
		if (-e "$1/primary") {
			$primary = read_file("$1/primary");
		}
		return 1 if $b eq "DIR::$1/$primary";
	}

	if ($b =~ /^DIR:([^:].*)$/) {
		if (-e "$1/primary") {
			$primary = read_file("$1/primary");
		}
		return 1 if $a eq "DIR::$1/$primary";
	}

	return 0;
}

sub ccache_is_default { cmp_ccnames(shift, $ccdefault); }

sub ccache_is_environ { cmp_ccnames(shift, $ccenviron); }

sub ccache_is_current { cmp_ccnames(shift, $cccurrent); }

sub put_env {
	my ($key, $val) = @_;
	$ENV{$key} = $val;

	for ($ENV{SHELL}) {
		if (m{/(sh|bash|zsh)$}) {
			$val =~ s/'/'\\''/g;
			say EVAL "$key=\'$val\'; export $key;";
		}
		else {
			_warn("unrecognized shell \"$ENV{SHELL}\"");
			say EVAL "$key=$val";
		}
	}
}

sub find_ccache_for_principal {
	my ($arg) = @_;

	my $max_expiry = 0;
	my $max_ccname;

	for my $ccname (@caches) {
		my $principal;
		my $ccrealm;
		my $expiry;
		my $tgt_expiry;
		my $init_expiry;

		$principal = read_proc("pklist", "-P", "-c", $ccname);
		if ($principal ne $arg) {
			next;
		}
		if ($principal =~ /.*@(.+)$/) {
			$ccrealm = $1;
		}

		open(my $proc, "-|", which("pklist"), "-c", $ccname)
		or _die("could not run 'pklist': $!");
		my (@fields, %row);
		while (my $line = <$proc>) {
			chomp($line);
			my @l = split(/\t/, $line);
			for (shift @l) {
				if ($_ eq "CREDENTIALS") {
					@fields = @l;
				}
				elsif ($_ eq "ticket") {
					_die("pklist output was missing header line") if !@l;
					@row{@fields} = @l;
					my $t_client = $row{client_name};
					my $t_service = $row{server_name};
					my $t_expiry = $row{expiry_time};
					my $t_flags = $row{flags};
					if ($t_service eq "krbtgt/$ccrealm\@$ccrealm") {
						$tgt_expiry = $t_expiry;
					}
					if ($t_flags =~ /I/) {
						$init_expiry = $t_expiry;
					}
				}
			}
		}
		close($proc);

		$expiry = $tgt_expiry || $init_expiry;
		if ($expiry > $max_expiry) {
			$max_expiry = $expiry;
			$max_ccname = $ccname;
		}
	}

	return ($max_ccname, $max_expiry);
}

sub switch_ccache {
	my ($ccname) = @_;

	return 0 if !$ccname;
	return 0 if !$can_switch;

	for ($ccname) {
		if (m|^DIR::(.+)$|) {
			my $ccdirname = "DIR:".dirname($1);
			put_env("KRB5CCNAME", $ccdirname);
			run_proc("kswitch", "-c", $ccname);
		}
		else {
			put_env("KRB5CCNAME", $ccname);
		}
	}

	if (run_proc("pklist", "-q") == 0) {
		my $princ = read_proc("pklist", "-P");
		say "Switched to \e[1m$princ\e[m ($ccname)";
	} else {
		say "New ccache ($ccname)";
	}

	return 1;
}

sub do_print_ccache {
	my ($ccname, $num) = @_;

	my $valid;
	my $shortname;
	my $principal;
	my $ccrealm;
	my $expiry;
	my $tgt_expiry;
	my $init_service;
	my $init_expiry;

	my $row_fmt;
	my $item_flag;
	my $expiry_str;

	my $num_tickets;

	$shortname = collapse_ccname($ccname);

	_debug("examining ccache '$ccname' aka '$shortname'");

	if (ccache_is_current($ccname)) {
		$item_flag = "‣";
	}

	$valid = run_proc("pklist", "-q", "-c", $ccname) == 0;
	if (!$valid) {
		$row_fmt = "empty";
		$principal = "(none)";
		$expiry_str = "(empty)";
		goto do_print;
	}

	open(my $proc, "-|", which("pklist"), "-c", $ccname)
	or _die("could not run 'pklist': $!");
	my (@fields, %row);
	while (<$proc>) {
		chomp;
		my @l = split(/\t/, $_);
		_debug("- pklist output: '@l'");
		for (shift @l) {
			if ($_ eq "principal") {
				($principal) = @l;
				if ($principal =~ /.*@(.+)$/) {
					$ccrealm = $1;
				}
			}
			elsif ($_ eq "CREDENTIALS") {
				@fields = @l;
			}
			elsif ($_ eq "ticket") {
				_die("pklist output was missing header line") if !@l;
				@row{@fields} = @l;
				my $t_client = $row{client_name};
				my $t_service = $row{server_name};
				my $t_expiry = $row{expiry_time};
				my $t_flags = $row{flags};

				if ($t_service eq "krbtgt/$ccrealm\@$ccrealm") {
					$tgt_expiry = $t_expiry;
				}
				if ($t_flags =~ /I/) {
					$init_service = $t_service;
					$init_expiry = $t_expiry;
				}
				++$num_tickets;
			}
		}
	}
	close($proc);

	if (!defined $principal) {
		_debug("no client principal in output, skipping ccache");
		return 0;
	}

	if (!$num_tickets) {
		$row_fmt = "empty";
		$expiry_str = "(empty)";
		goto do_print;
	}

	$expiry = $tgt_expiry || $init_expiry || 0;

	if ($expiry) {
		if ($expiry <= time) {
			$item_flag //= "×";
			$expiry_str = "expired";
		} else {
			_debug("- expires in ".($expiry - time)." seconds");
			$expiry_str = interval($expiry);
		}

		if ($expiry <= time) {
			$row_fmt = "expired";
		} elsif ($expiry <= time + 15*60) {
			$row_fmt = "expiring"
		} else {
			$row_fmt = "normal";
		}
	}

do_print:
	if (ccache_is_current($ccname)) {
		$row_fmt .= "_active";
	}

	_debugvar("row_fmt" => $row_fmt);

	my $row_color = $COLORS{$row_fmt};
	my $name_color = $COLORS{"name_".$row_fmt}
			|| $COLORS{$row_fmt eq "normal" ? "name" : $row_fmt};
	my $princ_color = $COLORS{"princ_".$row_fmt}
			|| $COLORS{$row_fmt eq "normal" ? "princ" : $row_fmt};
	my $time_color = $COLORS{"time_".$row_fmt}
			|| $COLORS{$row_fmt eq "normal" ? "time" : $row_fmt};

	_debugvar("init_service", $init_service);
	_debugvar("ccrealm", $ccrealm);

	if (defined $ccrealm && $ccrealm eq "WELLKNOWN:ANONYMOUS"
	    && $init_service =~ /^krbtgt\/.*@(.+)$/) {
		$ccrealm = $1;
		$principal = "\@$1 (anonymous)";
	}

	print fmt($row_color, $item_flag, "%1s");
	printf " %2d ", $num+1;
	print fmt($name_color, $shortname, "%-15s");
	if (length($shortname) > 15) {
		print "\n", " "x20;
	}
	print " ", fmt($princ_color, $principal, "%-40s");
	print " ", fmt($time_color, $expiry_str, "%8s");
	print "\n";

	if (defined $ccrealm && defined $init_service
	    && $init_service ne "krbtgt/".$ccrealm."@".$ccrealm) {
		print " "x20, " for ", $init_service, "\n";
	}

	return 1;
}

if (-t 1 && $ENV{TERM}) {
	%COLORS = map {$_ => "\e[$COLORS{$_}m"} keys %COLORS;
} else {
	%COLORS = map {$_ => ""} keys %COLORS;
}

if (!which("pklist")) {
	_die("'pklist' must be installed to use this tool");
}

open(EVAL, ">&=", 3) or do {
	_warn("cache switching unavailable (could not open fd#3)");
	$can_switch = 0;
	open(EVAL, ">/dev/null");
};

$rundir = $ENV{XDG_RUNTIME_DIR} || $ENV{XDG_CACHE_HOME} || $ENV{HOME}."/.cache";
$ccprefix = "/tmp/krb5cc_${UID}_";
$runprefix = "$rundir/krb5cc";

chomp($cccurrent = qx(pklist -N));
chomp($ccdefault = qx(unset KRB5CCNAME; pklist -N));
$ccenviron = $ENV{KRB5CCNAME} // "";

$cccdir = "";
$cccprimary = "";
if (-d $runprefix) {
	$cccdir = $runprefix;
}
if ($cccurrent =~ m|^DIR::(.+)$|) {
	$cccdir = dirname($1);
	if (-f "$cccdir/primary") {
		$cccprimary = read_file("$cccdir/primary");
	} else {
		$cccprimary = "tkt";
	}
}

@caches = enum_ccaches();

my $cmd = shift @ARGV;

for ($cmd) {
	if (!defined $_) {
		my $num = 0;
		for my $ccname (@caches) {
			$num += do_print_ccache($ccname, $num);
		}
		if (!$num) {
			say "No Kerberos credential caches found.";
			exit 1;
		}
	}
	elsif (/^--help$/) {
		usage();
		exit;
	}
	elsif ($_ eq "purge") {
		for my $ccname (@caches) {
			my $principal = read_proc("pklist", "-c", $ccname, "-P");
			say "Renewing credentials for $principal in $ccname";
			run_proc("kinit", "-c", $ccname, "-R") == 0
			|| run_proc("kdestroy", "-c", $ccname);
		}
	}
	elsif ($_ eq "destroy") {
		my @destroy = grep {defined} map {expand_ccname($_)} @ARGV;
		run_proc("kdestroy", "-c", $_) for @destroy;
	}
	elsif ($_ eq "clean") {
		say "Destroying all credential caches.";
		run_proc("kdestroy", "-c", $_) for @caches;
	}
	elsif ($_ eq "expand") {
		say expand_ccname($_) for @ARGV;
	}
	elsif ($_ eq "list") {
		say for @caches;
	}
	elsif ($_ eq "slist") {
		say collapse_ccname($_) for @caches;
	}
	elsif ($_ eq "trace") {
		$ENV{KRB5_TRACE} = "/dev/stderr";
		system {$ARGV[0]} @ARGV;
	}
	elsif ($_ eq "test-roundtrip") {
		for my $name (@caches) {
			my $tmp;
			say " original: ", ($tmp = $name);
			say "collapsed: ", ($tmp = collapse_ccname($tmp));
			say " expanded: ", ($tmp = expand_ccname($tmp));
			say "   result: ", ($tmp eq $name ? "\e[1;32mPASS\e[m"
			                                  : "\e[1;31mFAIL\e[m");
			say "";
		}
	}
	elsif ($_ eq "dump-aliases") {
		my %aliases = read_aliases();
		for (sort keys %aliases) {
			say $_."\t-> ".join(" ", @{$aliases{$_}});
		}
	}
	elsif (/^=(.*)$/) {
		my %aliases = read_aliases();
		my $alias = $aliases{$1};
		if (!$alias) {
			_die("alias '$1' not defined");
		}
		my $ccname = expand_ccname($1);
		switch_ccache($ccname);
		if (run_proc("klist", "-s") > 0) {
			exit run_proc("kinit", @$alias) >> 8;
		}
	}
	elsif (/.+@.+/) {
		my ($ccname, $expiry) = find_ccache_for_principal($cmd);
		if ($expiry) {
			switch_ccache($ccname) || exit 1;
		} else {
			switch_ccache("new") || exit 1;
			run_proc("kinit", $cmd, @ARGV);
		}
	}
	else {
		my $ccname = expand_ccname($cmd);
		if ($ccname) {
			switch_ccache($ccname) || exit 1;
			run_proc("kinit", @ARGV) if @ARGV;
		} else {
			_die("'$cmd' is neither a command nor ccache name");
		}
	}
}
