#!/usr/bin/perl

use strict;
use warnings;

use Expect;
use Carp;
use File::Spec;
use Regexp::Common qw( net );
use MIME::Base64;
use Getopt::Long;


#------------------------------------------------------------------------#
# User Parameter Parsing
my %OPTS = ();
GetOptions(\%OPTS, qw(
	define=s
));
#------------------------------------------------------------------------#

#------------------------------------------------------------------------#
$Expect::Log_Stdout = 0;
my %BIN = (
	manage_agents		=> '/var/ossec/bin/manage_agents',
	list_agents			=> '/var/ossec/bin/list_agents',
);
#------------------------------------------------------------------------#
# Config
my %CFG = (
	key_dir		=> '/usr/local/cfkeys/ossec',
	mode		=> exists $OPTS{define} && length $OPTS{define} ? 'cfengine' : 'normal',
); 
#------------------------------------------------------------------------#

#------------------------------------------------------------------------#
# Clean Data
#------------------------------------------------------------------------#
my $HOSTNAME = lc shift;
my $IP = shift;
my $KEY = undef;
my $ID = undef;
my $restartAgent = undef;

$HOSTNAME =~ s/[^\w\d\-_]//g;
$IP =~ s/[^\d\.]//g;

croak "usage: $0 hostname ip\n" unless length($HOSTNAME) && length($IP);

#------------------------------------------------------------------------#
# Check the server to see if the agent is listed.
my %agents = ();
my $menuSeen = 0;
my $exp = new Expect();
$exp->spawn( $BIN{manage_agents} );
$exp->expect(30,
	[
		qr{Choose your action: A,E,L,R or Q:\s+},
		sub {
			my $fh = shift;
			if( $menuSeen == 1 ) {
				$fh->send("Q\n");
			}
			else {
				$fh->send("L\n");
				$menuSeen = 1;
			}
			exp_continue;
		}
	],
	[
		qr{ID: (\d+), Name: ([^\,]+), IP: $RE{net}{IPv4}{-keep}},
		sub {
			my $fh = shift;
			my ($id, $hostname, $ip) = $exp->matchlist;
			$agents{$id} = { hostname => $hostname, ip => $ip };
			exp_continue;
		}
	],
	[
		qr{Press ENTER to return to the main menu},
		sub {
			my $fh = shift;
			$fh->send("\n");
			exp_continue;
		}
	],

);

my @REMOVE = ();
my $createKey = 1;
foreach my $id ( keys %agents ) {
	my $host = $agents{$id}->{hostname};
	my $ip = $agents{$id}->{ip};
	if( $host eq $HOSTNAME ) {
		if( $ip eq $IP ) {
			$createKey = 0;
			last;	
		}
		else {
			push @REMOVE, { id => $id, hostname => $host, ip  => $ip };
		}
	}
}

#
# Remove any keys that we need to remove:
foreach my $href (@REMOVE) {
	showMessage("!! Removing key for pair ( $href->{hostname}, $href->{ip} )\n");
	my $exp = new Expect;
	my $menuSeen = 0;
	$exp = new Expect;
	$exp->spawn( $BIN{manage_agents} );
	$exp->expect( 30,
		[
			qr{Choose your action: A,E,L,R or Q:\s+},
			sub {
				my $fh = shift;
				if( $menuSeen == 1 ) {
					$fh->send("Q\n");
				}
				else {
					$fh->send("R\n");
					$menuSeen = 1;
				}
				exp_continue;
			}
		],
		[
			qr{Provide the ID of the agent to be removed.*:\s+},
			sub {
				my $fh = shift;
				$fh->send("$href->{id}\n");
				exp_continue;
			}
		],
		[
			qr{Confirm deleting it\?\(y/n\):\s+},
			sub {
				my $fh = shift;
				$fh->send("y\n");
				exp_continue;
			}
		],

	);

}

if( $createKey == 1 ) {
	showMessage("++ Creating key for $HOSTNAME at $IP..\n");
	$restartAgent = 1;
	my $exp = new Expect;
	my $menuSeen = 0;
	$exp = new Expect;
	$exp->spawn( $BIN{manage_agents} );
	$exp->expect( 30,
		[
			qr{Choose your action: A,E,L,R or Q:\s+},
			sub {
				my $fh = shift;
				if( $menuSeen == 1 ) {
					$fh->send("Q\n");
				}
				else {
					$fh->send("A\n");
					$menuSeen = 1;
				}
				exp_continue;
			}
		],
		[
			qr{\* A name for the new agent:\s+},
			sub {
				my $fh = shift;
				$fh->send("$HOSTNAME\n");
				exp_continue;
			}
		],
		[
			qr{\* The IP Address of the new agent:\s+},
			sub {
				my $fh = shift;
				$fh->send("$IP\n");
				exp_continue;
			}
		],
		[
			qr{\* An ID for the new agent\[\d+\]:\s+},
			sub {
				my $fh = shift;
				$fh->send("\n");
				exp_continue;
			}
		],
		[
			qr{Confirm adding it\?\(y/n\):\s+},
			sub {
				my $fh = shift;
				$fh->send("y\n");
				exp_continue;
			}
		],

	);
}

#
# Extract the Key
$menuSeen = 0;
my $grabKey = 0;
$exp = new Expect;
$exp->spawn( $BIN{manage_agents} );
$exp->expect( 10,
	[
		qr{Choose your action: A,E,L,R or Q:\s+},
		sub {
			my $fh = shift;
			if( $menuSeen == 1 ) {
				$fh->send("Q\n");
			}
			else {
				$fh->send("E\n");
				$menuSeen = 1;
			}
			exp_continue;
		}
	],
	[
		qr{\s+ID: (\d+), Name: ([^\,]+), IP: $RE{net}{IPv4}{-keep}},
		sub {
			my $fh = shift;
			my ($id, $hostname, $ip) = $exp->matchlist;
			if( $hostname eq $HOSTNAME && $ip eq $IP ) {
				$ID = $id;
			}
			exp_continue;
		}
	],
	[
		qr{extract the key .*:\s+},
		sub {
			my $fh = shift;
			if( defined $ID ) {
				$fh->send("$ID\n");
			}
			else {
				$fh->send("\\q\n");
			}
			exp_continue;
		}
	],
	[
			qr{Agent key information},
			sub {
				$grabKey = 1;
				exp_continue;
			},
	],
	[
			qr{^(\S+)\s+$},
			sub {
				my $fh = shift;
				my ($string) = $fh->matchlist();
				if( $grabKey ) {
					$KEY = $string;
					$grabKey = 0;
				}
				exp_continue;
			},
	],
	[
			qr{Press ENTER to return},
			sub {
				my ($fh) = @_;
				$fh->send("\n");
				exp_continue;
			},
	],

);

#
# Extract the key string:
my $KEYSTRING = decode_base64($KEY);
#
# Write to the client key file
my $filename = File::Spec->catfile( $CFG{key_dir}, $HOSTNAME . '.ossec' );
open(my $fh, '>', $filename) 
	or croak "unable to write to $filename: $!\n";
print $fh $KEYSTRING;
close $fh;
showMessage("$KEYSTRING\n");

if( $restartAgent && $CFG{mode} eq 'cfengine' ) {
	print $OPTS{define},"\n";
}

sub showMessage {
	if( $CFG{mode} eq 'normal' ) {
		for(@_) {
			print $_;
		}
	}
}