#!/usr/bin/perl -w

use strict;

use Net::SNMP;
use Data::Dumper;

# Multiple PDU ips
# server name
# community
# ACTION

# script   server action community ip [ip...]

# action is one of:
#   on <node>
#   off <node>
#   reset <node>
#   status <node>
#   nodelist
#   ping


my $node = shift;
my $action = shift;
my $community = shift;
my @ips = @ARGV;

if ( scalar @ips < 1) {
	print STDERR "Usage: $0 node action community ip [ip...]\n";
	exit(1);
}

my %oid;
$oid{'generic'}->{'sysObjectID'} = '.1.3.6.1.2.1.1.2.0';

my $sysids;
$sysids->{'.1.3.6.1.4.1.318.1.3.4.5'}         = 'rackpdu';
$sysids->{'.1.3.6.1.4.1.318.1.3.4.5.1.3.4.5'} = 'rackpdu';
$sysids->{'.1.3.6.1.4.1.1718.3'}              = 'switchedcdu';
$sysids->{'.1.3.6.1.4.1.17420'}               = 'amazingpdu';


my $pduinfo;
$pduinfo->{'rackpdu'}->{'outlet_base'} = '.1.3.6.1.4.1.318.1.1.12.3.5.1.1.2';
$pduinfo->{'rackpdu'}->{'control_oid'} = '.1.3.6.1.4.1.318.1.1.12.3.3.1.1.4.%s';
$pduinfo->{'rackpdu'}->{'status_oid'} = '.1.3.6.1.4.1.318.1.1.12.3.5.1.1.4.%s';
$pduinfo->{'rackpdu'}->{'func'} = \&rackpdu_control;
$pduinfo->{'switchedcdu'}->{'outlet_base'} = '.1.3.6.1.4.1.1718.3.2.3.1.3.1';
$pduinfo->{'switchedcdu'}->{'control_oid'} = '.1.3.6.1.4.1.1718.3.2.3.1.11.1.%s';
$pduinfo->{'switchedcdu'}->{'status_oid'} = '.1.3.6.1.4.1.1718.3.2.3.1.5.1.%s';
$pduinfo->{'switchedcdu'}->{'func'} = \&switchedcdu_control;
$pduinfo->{'amazingpdu'}->{'outlet_base'} = '.1.3.6.1.4.1.17420.1.2.9.1.14';
$pduinfo->{'amazingpdu'}->{'control_oid'} = '.1.3.6.1.4.1.17420.1.2.9.1.13.0';
$pduinfo->{'amazingpdu'}->{'func'} = \&amazingpdu_control;



# Connect to all the PDUs
my ($conn, $ip, $type, $session, $func);

foreach $ip (@ips) {
	print STDERR "Connecting to $ip\n";
	my ($s, $t) = pduconnect($ip, $community);
	$conn->{$ip}->{'session'} = $s;
	if (!defined $sysids->{$t}) {
		print STDERR "Unknown PDU sysid: $t\n";
		exit(1);
	}
	$conn->{$ip}->{'type'} = $sysids->{$t};
}

if ($action eq 'ping') {
	# Assume OK, since we have just read the sysObjectID
	exit 0;
}

# Read info from PDUs
# $conn -> 
#   $ip ->
#     type = pdu type
#     session = snmp session object
#     func = function to get/set info 
#     nodelist = hash of:  {node name} = [ port nos ]
#     portlist = hash of:  {port no} -> node name

foreach $ip (keys %$conn) {
	print STDERR "Reading port list from $ip\n";
	$type    = $conn->{$ip}->{'type'};
	$session = $conn->{$ip}->{'session'};
	$func    = $pduinfo->{$type}->{'func'};
	
	my ($nodelist, $portlist) = &$func('list', $session);
	
	$conn->{$ip}->{'portlist'} = $portlist;
	$conn->{$ip}->{'nodelist'} = $nodelist;
}

if ($action eq 'nodelist') {
	my $hash;
   foreach $ip (keys %$conn) {
		my $list = $conn->{$ip}->{'nodelist'};
		if (defined $list) {
			foreach my $item (keys %$list) {
				$hash->{$item} = 1;
			}
		}
	}
	foreach my $name (sort keys %$hash) {
		print "$name\n";
	}
	exit;
}


if ($action eq 'status') {
	my $RV = 1;
	foreach $ip (keys %$conn) {
		$type    = $conn->{$ip}->{'type'};
		$session = $conn->{$ip}->{'session'};
		$func    = $pduinfo->{$type}->{'func'};
		
		my $portlist = $conn->{$ip}->{'nodelist'}->{$node};
		if (!defined $portlist) {
#			print STDERR "Cannot get status - node $node is not managed here\n";
			print STDERR "$ip: Skipping status - node $node is not managed here\n";
			next;
#			exit 2;
		}
		foreach my $port (@$portlist) {
			print STDERR "Checking $port on $ip ... ";
			my $stat = &$func('status', $session, $port);
			print STDERR "$stat\n";
			$RV = 0 if $stat eq 'on';
		}	
	}
	exit $RV;
}


# Try turn off on each PDU
# Confirm on each PDU
if ($action eq 'off' or $action eq 'reset') {
	foreach $ip (keys %$conn) {
		$type    = $conn->{$ip}->{'type'};
		$session = $conn->{$ip}->{'session'};
		$func    = $pduinfo->{$type}->{'func'};
		
		my $portlist = $conn->{$ip}->{'nodelist'}->{$node};
		if (!defined $portlist) {
#			print STDERR "Cannot turn off - node $node is not managed here\n";
			print STDERR "$ip: Skipping turn off - node $node is not managed here\n";
			next;
#			exit 1;
		}
		
		foreach my $port (@$portlist) {
			print STDERR "Turning off port $port on $ip\n";
			&$func('off', $session, $port);
		}	
		sleep 1;
		foreach my $port (@$portlist) {
			print STDERR "Checking $port on $ip ... ";
			my $stat = &$func('status', $session, $port);
			print STDERR "$stat\n";

			if ($stat ne 'off') {
				print STDERR "\n";
				print STDERR "Port $port on $ip did not turn off\n";
				exit(1);
			}
		}	
	}
}

# Try turn on on each PDU
# Confirm on each PDU
if ($action eq 'on' or $action eq 'reset') {
	foreach $ip (keys %$conn) {
		$type    = $conn->{$ip}->{'type'};
		$session = $conn->{$ip}->{'session'};
		$func    = $pduinfo->{$type}->{'func'};
		
		my $portlist = $conn->{$ip}->{'nodelist'}->{$node};
		if (!defined $portlist) {
#			print STDERR "Cannot turn on - node $node is not managed here\n";
			print STDERR "$ip: Skipping turn on - node $node is not managed here\n";
			next;
#			exit 1;
		}
		
		foreach my $port (@$portlist) {
			print STDERR "Turning on port $port on $ip\n";
			&$func('on', $session, $port);
		}	
		sleep 1;
		foreach my $port (@$portlist) {
			print STDERR "Checking $port on $ip ... ";
			my $stat = &$func('status', $session, $port);
			print STDERR "$stat\n";

			if ($stat ne 'on') {
				print STDERR "\n";
				print STDERR "Port $port on $ip did not turn on\n";
				exit(1);
			}
		}	
	}
}



#my $outlets = walk($oid->{'portlistbase'});
#print Dumper($outlets);

#foreach my $port ( @{$outlets->{'ports'}->{$host}}) {
	#set($oid->{'portcontrolbase'}.".".$port, INTEGER, $options->{$option});
#}

#my $status = walk($oid->{'portstatusbase'});
#print Dumper($status);


sub pduconnect {
	my $ip = shift;
	my $community = shift;

	my ($session, $error) = Net::SNMP->session(
		-hostname  => $ip,
		-community => $community,
		-port      => 161,
	);

	if (!defined $session) {
		print STDERR "ERROR: IP=$ip $error\n";
		exit 1;
	}

	# Read the PDU type:
	my $oid = $oid{'generic'}->{'sysObjectID'};
	my $sysid = get($session, $oid);


	return ($session, $sysid);
}

sub get {
	my $session = shift;
	my $o = shift;
	my $type = shift;
	my $val = shift;
	
	my $result;
	
	$result = $session->get_request(
		-varbindlist => [ $o ]
	);

	if (!defined $result) {
		print STDERR "ERROR get $o: ".$session->error()."\n";
		exit 1;
	}
	
	return $result->{$o};
}

sub set {
	my $session = shift;
	my $o = shift;
	my $type = shift;
	my $val = shift;
	
	my $result;
	
	$result = $session->set_request(
		-varbindlist => [ $o, $type, $val ]
	);

	if (!defined $result) {
		print STDERR "ERROR set: ".$session->error()."\n";
		exit 1;
	}
	
}


sub walk {
	my $session = shift;
	my $o = shift;
	my $result;
	
	my $rv;

	my $inoid = $o;
	while ($result = $session->get_next_request(
		-varbindlist => [ $inoid ],
	)) {
		$inoid = (keys %$result)[0];
		last if substr($inoid, 0, length($o)) ne $o;
		my $name = $result->{$inoid};
		$rv->{$inoid} = $name;
	}

	if (!defined $result) {
		print STDERR "ERROR: ".$session->error()."\n";
		exit 1;
	}
	
	return $rv;
}

#     nodelist = hash of:  {node name} = [ port nos ]
#     portlist = hash of:  {port no} -> node name
sub rackpdu_control {
	my $action = shift;
	my $session = shift;


	if ($action eq 'list') {
		my $oid = $pduinfo->{'rackpdu'}->{'outlet_base'};
		my $list = walk($session, $oid);
		my $nl;
		my $pl;

		foreach my $o (keys %$list) {
			my $name = $list->{$o};

			$o =~ /^.*\.(\d+)$/;
			my $port = $1;
			push @{$nl->{$name}}, $port;
			$pl->{$port} = $name;
		}	
		return ($nl, $pl);
	}
	
	if ($action eq 'off') {
		my $port = shift;
		my $oid = $pduinfo->{'rackpdu'}->{'control_oid'};
		$oid = sprintf($oid, $port);
		set($session, $oid, INTEGER, 2);
		return;
	}
	
	if ($action eq 'on') {
		my $port = shift;
		my $oid = $pduinfo->{'rackpdu'}->{'control_oid'};
		$oid = sprintf($oid, $port);
		set($session, $oid, INTEGER, 1);
		return;
	}
	
	if ($action eq 'status') {
		my $port = shift;
		my $oid = $pduinfo->{'rackpdu'}->{'status_oid'};
		$oid = sprintf($oid, $port);
		my $stat = get($session, $oid);
		return 'off' if $stat eq '2';
		return 'on' if $stat eq '1';
		return 'unknown';
	}

	print STDERR "Unknown action: rackpdu: $action\n";
	exit(1);	
}

#     nodelist = hash of:  {node name} = [ port nos ]
#     portlist = hash of:  {port no} -> node name
sub switchedcdu_control {
	my $action = shift;
	my $session = shift;


	if ($action eq 'list') {
		my $oid = $pduinfo->{'switchedcdu'}->{'outlet_base'};
		my $list = walk($session, $oid);
		my $nl;
		my $pl;

		foreach my $o (keys %$list) {
			my $name = $list->{$o};
			$name =~ s/:.*//g;

			$o =~ /^.*\.(\d+\.\d+)$/;
			my $port = $1;
			push @{$nl->{$name}}, $port;
			$pl->{$port} = $name;
		}	
		return ($nl, $pl);
	}
	
	if ($action eq 'off') {
		my $port = shift;
		my $oid = $pduinfo->{'switchedcdu'}->{'control_oid'};
		$oid = sprintf($oid, $port);
		set($session, $oid, INTEGER, 2);
		return;
	}
	
	if ($action eq 'on') {
		my $port = shift;
		my $oid = $pduinfo->{'switchedcdu'}->{'control_oid'};
		$oid = sprintf($oid, $port);
		set($session, $oid, INTEGER, 1);
		return;
	}
	
	if ($action eq 'status') {
		my $port = shift;
		my $oid = $pduinfo->{'switchedcdu'}->{'status_oid'};
		$oid = sprintf($oid, $port);
		my $stat = get($session, $oid);
		return 'off' if $stat eq '0';
		return 'on' if $stat eq '1';
		return 'unknown';
	}

	print STDERR "Unknown action: switchedcdu: $action\n";
	exit(1);	
}

sub amazingpdu_control {
	my $action = shift;
	my $session = shift;
	if ($action eq 'list') {
		my $oid = $pduinfo->{'amazingpdu'}->{'outlet_base'};
		my $list = walk($session, $oid);
		my $nl;
		my $pl;

		foreach my $o (keys %$list) {
			my $name = $list->{$o};
			$name =~ s/,.*//g;

			$o =~ /^.*\.(\d+)\.0$/;
			my $port = $1;
			push @{$nl->{$name}}, $port;
			$pl->{$port} = $name;
		}	
		return ($nl, $pl);
	}

	if ($action eq 'off') {
		my $port = shift;
		my $oid = $pduinfo->{'amazingpdu'}->{'control_oid'};
		my @current = split(/,/, get($session, $oid));
		$current[$port-1] = 0;
		set($session, $oid, OCTET_STRING, join(",", @current));
		return;
	}
	
	if ($action eq 'on') {
		my $port = shift;
		my $oid = $pduinfo->{'amazingpdu'}->{'control_oid'};
		my @current = split(/,/, get($session, $oid));
		$current[$port-1] = 1;
		set($session, $oid, OCTET_STRING, join(",", @current));
		return;
	}
	
	if ($action eq 'status') {
		my $port = shift;
		my $oid = $pduinfo->{'amazingpdu'}->{'control_oid'};
		my @current = split(/,/, get($session, $oid));
		my $stat = $current[$port-1];
		return 'off' if $stat eq '0';
		return 'on' if $stat eq '1';
		return 'unknown';
	}

	print STDERR "Unknown action: amazingpdu: $action\n";
	exit(1);	

}