opentrafficshaperd

#!/usr/bin/perl
# Main OpenTrafficShaper program
# Copyright (C) 2007-2014, AllWorldIT
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

use strict;
use warnings;


# Set the dirs we look for library files in
use lib(
	'/usr/local/lib/opentrafficshaper-1.0',
	'/usr/lib/opentrafficshaper-1.0',
	'/usr/lib64/opentrafficshaper-1.0',
	'opentrafficshaper',
	'awitpt/lib',
);

# Enable assertions
BEGIN {
	package POE::Kernel;
	use constant ASSERT_DEFAULT => 1;
}

# System stuff we need
use Config::IniFiles;
use Getopt::Long;
use POE;
use POSIX qw(setsid);
use Time::HiRes qw(time);

# Our own stuff
use opentrafficshaper::version;
use opentrafficshaper::logger;
use opentrafficshaper::plugins qw( plugin_register );


# Main config
my $globals;
# We just create the logger first, its only using STDERR here
my $logger = new opentrafficshaper::logger;


#
# MAIN
#
parseCfgCmdLine();

# Check if we must use a log file instead
if (defined($globals->{'config'}->{'log_file'})) {
	$logger->open($globals->{'config'}->{'log_file'});
}
$logger->setLevel($globals->{'config'}->{'log_level'});
# Check if we need to go background
if (defined($globals->{'config'}->{'background'})) {
	daemonize();
}

displayBanner();
init();
start();
$logger->log(LOG_NOTICE,"[MAIN] Entering RUNNING state");
POE::Kernel->run();
exit;




# Function to display banner
sub displayBanner
{
	$logger->log(LOG_NOTICE,"[MAIN] OpenTrafficShaper v%s - Copyright (c) 2007-2014, AllWorldIT",VERSION);
}


# Function to parse our config and commandline
sub parseCfgCmdLine
{

	# Set defaults
	my $cfg;
	$cfg->{'config_file'} = "/etc/opentrafficshaper.conf";

	$cfg->{'background'} = "yes";
	$cfg->{'pid_file'} = "/var/run/opentrafficshaper/opentrafficshaperd.pid";
	$cfg->{'log_level'} = 2;
	$cfg->{'log_file'} = "/var/log/opentrafficshaper/opentrafficshaperd.log";

	# Parse command line params
	my $cmdline;
	%{$cmdline} = ();
	GetOptions(
			\%{$cmdline},
			"help",
			"config:s",
			"debug",
			"fg",
	) or die "Error parsing commandline arguments";

	# Check for some args
	if ($cmdline->{'help'}) {
		displayHelp();
		exit 0;
	}
	if (defined($cmdline->{'config'}) && $cmdline->{'config'} ne "") {
		$cfg->{'config_file'} = $cmdline->{'config'};
	}

	# Check config file exists
	if (! -f $cfg->{'config_file'}) {
		die("No configuration file '".$cfg->{'config_file'}."' found!");
	}

	# Use config file, ignore case
	tie my %inifile, 'Config::IniFiles', (
			-file => $cfg->{'config_file'},
			-nocase => 1
	) or die "Failed to open config file '".$cfg->{'config_file'}."': $!";
	my $inifileHandle = tied( %inifile );
	# Copy config
	my %config = %inifile;

	# Pull in params for the server
	my @server_params = (
			'log_level','log_file',
			'pid_file',
			'background',
	);
	foreach my $param (@server_params) {
		$cfg->{$param} = $config{'system'}{$param} if (defined($config{'system'}{$param}));
	}

	# Override
	if ($cmdline->{'debug'}) {
		$cfg->{'log_level'} = 4;
		$cfg->{'debug'} = 1;
	}

	# If we set on commandline for foreground, keep in foreground
	if ($cmdline->{'fg'} || (defined($config{'system'}{'background'}) && $config{'system'}{'background'} eq "no" )) {
		$cfg->{'background'} = undef;
		$cfg->{'log_file'} = undef;
		$cfg->{'pid_file'} = undef;
	}

	#
	# System plugins
	#
	if (ref($config{'plugins'}{'load'}) eq "ARRAY") {
		foreach my $plugin (@{$config{'plugins'}{'load'}}) {
			$plugin =~ s/\s+//g;
	 		# Skip comments
	 		next if ($plugin =~ /^#/);
			push(@{$cfg->{'plugin_list'}},$plugin);
		}
	} elsif (defined($config{'plugins'}{'load'})) {
		my @pluginList = split(/\s+/,$config{'plugins'}{'load'});
		foreach my $plugin (@pluginList) {
			# Skip comments
			next if ($plugin =~ /^#/);
			push(@{$cfg->{'plugin_list'}},$plugin);
		}
	}

	# We may have config file groups we want to remember for other plugins
	foreach my $group ($inifileHandle->Groups()) {
		# Loop with group members
		foreach my $member ($inifileHandle->GroupMembers($group)) {
			# Chop off group name and just get the member
			my $cleanMember = substr($member,length($group)+1);
			# Link the config...
			$config{$group}->{$cleanMember} = $config{$member};
		}
	}

	$globals->{'file.config'} = \%config;
	$globals->{'config'} = $cfg;
}


# Display help
sub displayHelp {
	displayBanner();

	print(STDERR<<EOF);

Usage: $0 [args]
    --config=<file>        Configuration file
    --debug                Put into debug mode
    --fg                   Don't go into background

EOF
}


# Initialize things we need
sub init
{
	# Certain things we need
	$globals->{'users'} = { };
	$globals->{'logger'} = $logger;
	$globals->{'version'} = VERSION;

	$logger->log(LOG_NOTICE,"[MAIN] Entering INITIALIZATION state");

	# Setup master session
	POE::Session->create(
		inline_states => {
			'_start' => \&main_session_start,
			'_stop' => \&main_session_stop,
			'main_SIGHUP' => \&main_signal_SIGHUP,
			'main_SIGINT' => \&main_signal_SIGINT,
		},
	);

	$logger->log(LOG_INFO,"[MAIN] Initializing plugins...");
	# We need to set the plugins global variable
	opentrafficshaper::plugins::init($globals);

	# Core configuration manager
	plugin_register("configmanager",1);

	# Load plugins
	foreach my $pluginName (@{$globals->{'config'}->{'plugin_list'}}) {
		plugin_register($pluginName,0);
	}

	$logger->log(LOG_INFO,"[MAIN] Plugins initialized");
}

# Function to start things up
sub start
{
	$logger->log(LOG_NOTICE,"[MAIN] Entering STARTING state");

	# Loop with plugins and call the start function for those that exist
	foreach my $pluginName (keys %{$globals->{'plugins'}})
	{
		my $plugin = $globals->{'plugins'}->{$pluginName};

		# Load the function up
		my $callStart = $plugin->{'Start'};
		if (defined($callStart)) {
			$callStart->();
		}
	}
}



# Become daemon
sub daemonize {
	chdir '/'
			or die "Can't chdir to /: $!";

	open STDIN, '/dev/null'
			or die "Can't read /dev/null: $!";

	open STDOUT, '> /dev/null'
			or die "Can't open stdout log: $!";

	defined(my $pid = fork)
			or die "Can't fork: $!";

	exit if $pid;

	# Write out our PID if we have a file to do it
	if (defined($globals->{'config'}->{'pid_file'})) {
		if (open(FH,"> ".$globals->{'config'}->{'pid_file'})) {
			print(FH $$);
			close(FH);
		} else {
			$logger->log(LOG_WARN,"[MAIN] Unable to write PID to '%s': %s",$globals->{'config'}->{'pid_file'},$!);
		}
	}

	setsid
			or die "Can't start a new session: $!";

	open STDERR, '> /dev/null'
			or die "Can't open stderr log: $!";
}



# Function to fire up our main session
sub main_session_start
{
	my $kernel = $_[KERNEL];

	$kernel->alias_set('main');

	# Register signal handlers
	$kernel->sig('HUP', 'main_SIGHUP');
	$kernel->sig('INT', 'main_SIGINT');
}


# Function to dispose of anything as a final stage to shutting down
sub main_session_stop
{
	$logger->log(LOG_DEBUG,"[MAIN] Shutdown");
}


# Function to handle SIGHUP
sub main_signal_SIGHUP
{
	my ($kernel,$signal_name) = @_[KERNEL,ARG0];

	$logger->log(LOG_NOTICE,"[MAIN] Got SIGHUP");

	$kernel->sig_handled();
}


# Function to handle SIGINT
sub main_signal_SIGINT
{
	my ($kernel, $signal_name) = @_[KERNEL, ARG0];

	$logger->log(LOG_NOTICE,"[MAIN] Got SIGINT, shutting down...");

	$kernel->signal($kernel,'handle_SIGINT');
}


# vim: ts=4