#!/usr/bin/perl
# quosnmp
# Version 1.0.0, Last modified on 2007-06-22
# A CUPS backend for SNMP-based print accounting and quota enforcement.
#
# Released by Marcus Lauer (marcus.lauer at nyu dot edu)
#
# Copyright (C) 2007 by Marcus Lauer (marcus.lauer at nyu dot edu) except where 
#   previous copyright is in effect.
#
# Based on accsnmp v1.02.20070124 by jeff hardy (hardyjm at potsdam dot edu)
# ############
# accsnmp
# v1.02.20070124
# jeff hardy (hardyjm at potsdam dot edu)
# backend wrapper hardware accounting for cups
#
# ############
# Copyright 2007, Jeff Hardy (hardyjm at potsdam dot edu)
#
# 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 2 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, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307,
# USA.
#
# ###########
# example URI - quosnmp://socket://192.168.xxx.xxx
#
# arguments according to the docs, as C sees them:
# 0     1     2         3        4      5            6
# lpd   jobid username  jobtitle copies printoptions file ///or read from STDIN if no filename
#
# arguments as Perl sees them:
# 0     1	 2        3      4            5
# jobid username jobtitle copies printoptions file ///or read from STDIN if no filename


use strict;
use Net::SNMP;

# --- Modify these settings to suit your needs.
my $ENFORCE_QUOTA = 1;					# Turn on/off [1/0] print quota enforcement
my $SAVE_JOBS = 0;					#turn on/off [1/0] saving printjobs
my $JOB_BLACKLIST = 0;					#turn on/off [1/0] jobs blacklist
my @JobBlacklist = ("");				# Names of jobs to be blacklisted.

my $LOGGING = 1;                                        # Turn on/off [1/0] print logging function call
my $LOGGING_FILE = "/var/log/cups/acc_page_log";        # File for printer logging
my $ERROR_LOGGING = 1;                                  # Turn on/off [1/0] error logging function call
my $ERROR_LOGGING_FILE = "/var/log/cups/acc_error_log"; # File for error logging

my $headerDiscount = 0;                                 # Credit accounts this many pages per print job for mandatory header (or footer) pages

my $DEBUG = 0;						#turn on/off [1/0] debugging function call
my $DEBUG_FILE = "/tmp/quosnmp.debug";			#tmp file for debugging

my $SNMP_COMMUNITY = "public";				#best to use read community obviously

# --- These settings can be modified, but it is less likely that you will need to do so.
my $accDir = "/var/log/cups/accounting";		#directory to keep accounting files
my $quoDir = "/var/log/cups/quotas";			#directory to keep quota files
my $accPrinterList = "/etc/cups/accounted-printers.txt";#list of special (usually color) printers -- pages printed to these count triple

my $backendDir = "/usr/lib/cups/backend";		#directory which contains all cups backends

my $SNMP_TIMEOUT = 15;                                  # SNMP timeout in seconds, 0 to try forever, max 60
my $SNMP_RETRIES = 5;					# Number of times to retry an SNMP request, max 20
my $BACKEND_ATTEMPTS = 10;				#backend retry attempts, 0 to try forever
my $STALL_TIMEOUT = 600;				# Assume printer is stalled after this many seconds, 0 to disable

# --- You almost certainly do not want to modify these variables.
my $PAGECOUNT_OID = "1.3.6.1.2.1.43.10.2.1.4.1.1";	#printer lifetime pagecount
my $PRINTERSTATUS_OID = "1.3.6.1.2.1.25.3.5.1.1.1";	#printer status


# --- Only experts should modify the code below this line.

# This is used by some of the logging features.
my @startTime = localtime(time);

### MAIN
{
	### ARGS CHECK
	if (!$ARGV[0]){	# Device discovery mode
		print ("network quosnmp \"Unknown\" \"Accounted Printer (SNMP)\"\n");
		exit 0;
	}
	if (scalar(@ARGV) < 5 || scalar(@ARGV) > 6){	# Error
		print STDERR ("ERROR: Usage: quosnmp job-id user title copies options [file]\n");
		exit 1;
	}
	# Now that we got this far, let us name the arguments to keep our sanity
	my ($jobID,$userName,$jobTitle,$copies,$printOptions,$printFile) = @ARGV;

	### ENV CHECK
	# URI parsing must fit syntax or die
	# Ex: quosnmp://lpd://192.168.xxx.xxx
        # Also supported: quosnmp://hp:/net/HP_LaserJet_xxxx_Series?ip=192.168.xxx.xxx
	$ENV{DEVICE_URI} =~ m#(\S+)://(\S+):/\S+?(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\S?# or 
		cleanExit($jobID,$ENV{PRINTER},"ERROR: URI must be in format: quosnmp://<original uri>\n");
	my $wrapBackend = $1;
	my $realBackend = $2;
	my $printerIP = $3;

	### BLACKLIST/QUOTA CHECKS
	if ( length($userName) < 1 ){	# No username
		cleanExit($jobID,$ENV{PRINTER},"ERROR: No username for printjob: user must authenticate\n");
	}

	if ($ENFORCE_QUOTA == 1) {
		my $printQuota = getPrintQuota($quoDir,$userName);

		if ( $printQuota != -1 ) { # -1 is a special value meaning unlimited printing is allowed.
			my $userPagesPrinted = getUserTotal("$accDir/$userName");

			# If the user is over quota, or is disallowed to print, or no quota file was found at all, then exit.
			if ( $userPagesPrinted >= $printQuota || $printQuota == 0 || $printQuota == -2 ) {
				cleanExit($jobID,$ENV{PRINTER},"ERROR: Account ($userName) is at or over quota ($printQuota pages) and cannot print\n");
			}
		}
	}

	if ($JOB_BLACKLIST){
		my $jobStatus = chkJobBlacklist($jobTitle);
		if ($jobStatus == -1){	# Job is blacklisted
			cleanExit($jobID,$ENV{PRINTER},"ERROR: Jobname ($jobTitle) is blacklisted from printing\n");
		}
	}
	
	
	### SPOOL JOB TO TEMP FILE
	my $jid = $jobID;
	my $uid = $userName;
	$jid =~ s/\W//g;	#sanity check
	$uid =~ s/\W//g;	#sanity check
	my $tempFile = "$ENV{TMPDIR}/$jid-$uid-cupsjob$$";

	if ($printFile){	# Only spool given print file if we want to save
		if ($SAVE_JOBS){
			open (OUT, ">$tempFile") or cleanExit($jid,$ENV{PRINTER},"ERROR: Cannot write $tempFile: $!\n");
			open (FH, "$printFile") or cleanExit($jid,$ENV{PRINTER},"ERROR: Cannot read $printFile: $!\n");
			while(<FH>){
				print OUT "$_";
			}
			close FH;
			close OUT;
		}
	}
	else { # Have to spool to temp if STDIN
		open (OUT, ">$tempFile") or cleanExit($jid,$ENV{PRINTER},"ERROR: Cannot write $tempFile: $!\n");
		while(<STDIN>){
			print OUT "$_";
		}
		close OUT;
	}
	
	
	### PRINTING
	# SNMP pagecount: first snmp query
	my $prePageCount = snmpGet($printerIP,$PAGECOUNT_OID,$jid);


	# NOTE: the next several lines are the guts of wrapping a cups backend
	# 1) Pull the device URI from the environment and fix it by removing the wrapper
	# 2) Tack this onto the front of the arg array to pass to the real backend
	# Note: 1 and 2 are required as different backends find the device URI by
	# different means.  LPD gets it from the arg, whereas IPP gets it from env.
	# 3) If the job arrived via STDIN, tack printFile onto the end of the arg array
	# This will happen on non-raw queues and needs to be checked, otherwise we
	# will dos the printer with a "zero byte" error.
	# 4) Call the real backend with this new argument array
	
	my $deviceURI = $ENV{DEVICE_URI};
	$deviceURI =~ s#$wrapBackend://##;
	$ENV{DEVICE_URI} = $deviceURI;

	if (!($printFile)){
		$copies = 1;
		$printFile = $tempFile;
	}
	my @argvNew = ($deviceURI,$jobID,$userName,$jobTitle,$copies,$printOptions,$printFile);
	
	# Loop until $attempts trying the backend
	my $attempts = $BACKEND_ATTEMPTS;
	my $exitval = 1;
	while ($exitval){
		# Override arg0
		$exitval = system {"$backendDir/$realBackend"} @argvNew;
		$exitval >>= 8;
		if ($attempts > 0){
			$attempts--;
			if ($attempts == 0){
				unlink $tempFile;
				cleanExit($jid,$ENV{PRINTER},"ERROR: Backend failed!\n");
			}
		}
	}
	
	# Now the real backend is processing the job, and we loop to keep an eye 
	# on the printer status until it is done.
	sleep 1;
	my $waitCount = 0;
	my $somethingPrinted = 0;
	$SIG{ALRM} = sub { cleanExit($jid,$ENV{PRINTER},"ERROR: Printer is stalled!\n") };
	while (1){
		alarm ($STALL_TIMEOUT);

		# 3 is idle, 4 is printing, 5 is warmup, 1 is other
		my $prStatus = snmpGet($printerIP,$PRINTERSTATUS_OID,$jid);
		if ($prStatus == 4){
			print STDERR "INFO: Printer status: printing\n";
			$somethingPrinted = 1;
		}
		elsif (($prStatus != 4) && ($somethingPrinted)){
			$waitCount += 1;
			# Wait until status is idle for several counts
			if ($waitCount == 3){
				print STDERR "INFO: Job printed successfully\n";
				last;
			}
		}
		else{
			print STDERR "INFO: Printer status: other\n";
		}
		sleep 2;
	}
	
	# SNMP pagecount: last snmp query
	my $postPageCount = snmpGet($printerIP,$PAGECOUNT_OID,$jid);
	
	
	### ACCOUNTING
	# Delta of our pagecounts
	my $totalPages = $postPageCount - $prePageCount;
	# Give a discount for header pages
	if ($headerDiscount > 0){
	        $totalPages = $totalPages - $headerDiscount;
	        if ($totalPages < 1){
	                $totalPages = 0;
	        }
	}
	# If no error from real backend (non-zero is error)...
	if (!$exitval){
		accounting($userName,$totalPages);
	}
	
	
	### CLEANUP/EXIT
	# If requested, write out a system log.
	if ( $LOGGING ) {
	        open (LOG_FH, ">>$LOGGING_FILE") or $LOGGING = 0;
	        flock LOG_FH, 2;
	        seek LOG_FH, 0 , 2;
	        printf LOG_FH "%04d-%02d-%02d\t%02d:%02d:%02d\t",$startTime[5]+1900,$startTime[4]+1,$startTime[3],$startTime[2],$startTime[1],$startTime[0];
	        print LOG_FH "$printerIP\t$jobID\t$userName\t$jobTitle\t$copies\t$prePageCount\t$postPageCount\t$exitval\t$totalPages\t$headerDiscount\n";
	        flock LOG_FH, 8;
	        close (LOG_FH);
	}
	if ($DEBUG == 1){
		#print STDERR ("ERROR: $postPageCount minus $prePageCount equals $totalPages\n");
		debug($DEBUG_FILE);
	}
	if (!$SAVE_JOBS){
		unlink $tempFile;
	}
	exit $exitval;
}


sub cleanExit{
	# Args: jobid, queuename, error message to print
	# Returns: nothing, deletes job and re-enables queue
	my $jobid = shift;
	my $queue = shift;
	my $errorMsg = shift;

	if ( length($errorMsg) > 0 ) {
		print STDERR $errorMsg;

		if ( $ERROR_LOGGING == 1 ) {
		       	my @errorTime = localtime(time);

		        open (ERR_LOG_FH, ">>$ERROR_LOGGING_FILE") or $ERROR_LOGGING = 0;
		        flock ERR_LOG_FH, 2;
		        seek ERR_LOG_FH, 0 , 2;
		        printf LOG_FH "%04d-%02d-%02d\t%02d:%02d:%02d\t",$errorTime[5]+1900,$errorTime[4]+1,$errorTime[3],$errorTime[2],$errorTime[1],$errorTime[0];
		        print ERR_LOG_FH "$jobid\t$queue\t$errorMsg\n";
		        flock ERR_LOG_FH, 8;
		        close (ERR_LOG_FH);
		}
	}
	sleep 1;
	system("/usr/bin/lprm","$jobid");
	sleep 1;
	system("/usr/sbin/cupsenable","$queue");
	sleep 1;
	exit 1;
}


sub snmpGet{
	# Args: ip, oid, jid
	# Returns: value of SNMP query. JobID needed only for exiting script cleanly if this fails.
	my $ip = shift;
	my $oid = shift;
	my $jid = shift;

	# Creating the session should always succeed
	my ($session, $error) = Net::SNMP->session(
		-hostname	=> $ip,
		-version	=> 'snmpv1',	#or 'snmpv2c' or 'snmpv3' as required
		-community	=> $SNMP_COMMUNITY,
		-timeout	=> $SNMP_TIMEOUT,		#default, max: 60
		-retries	=> $SNMP_RETRIES		#default, max: 20
	);
	# Paranoid check
	if (!defined($session)) {
		cleanExit($jid,$ENV{PRINTER},"ERROR: SNMP session creation error: $error\n");
	}

	# The get request will loop until the alarm is tripped
	my $result;
	my $err;

	eval{
		while(!defined($result)){
			$result = $session->get_request(
        		-varbindlist => [$oid]
	        	);
	        $err = $session->error;
		}
		alarm (0);	# cancel the alarm
	};
	if ($@ =~ /SNMP ERROR/){
		cleanExit($jid,$ENV{PRINTER},"ERROR: SNMP Error: $err\n");
	}
	# Paranoid check.  This is spurious because the eval loop won't let this happen.
	if (!defined($result)) {
		my $err = $session->error;
		$session->close;
		cleanExit($jid,$ENV{PRINTER},"ERROR: SNMP Error: $err\n");
	}

	my $return = $result->{$oid};
	$session->close;
	return $return;
}


sub chkJobBlacklist{
	# Args: jobname
	# Returns: -1/bad | 1/good
	# This is not terribly safe, but can be terribly handy
	my $jobName = shift;

	foreach my $blackJob (@JobBlacklist) {
		if ($blackJob =~ $jobName) { return -1; }
	}

	return 1;
}


sub accounting{
	# Args: user, totalpages, printername
	# Returns: nothing, logs to page_log and writes to user file in accounting directory
	my $user = shift;
	my $newPages = shift;	#what was printed this time

	if ($user eq ""){
		$user = "NO_USER";
	}
	my $accFile = "$accDir/$user";
                                                                                            
	# If color printer from our printer accounting list, multiply pagecount by three
	#
	# Example accPrinterList:
	# Every queuename ends with colon, color queues add =COLOR= 
	#
	# SAT104CL:=COLOR=
	# SAT325:
	# SIS217:
	#
	if (chkPrinterColor($ENV{PRINTER},$accPrinterList)){
        	$newPages = $newPages*3;
	}

	# Send nonstandard message to page_log (real wrapped backend will also send standard message)
	print STDERR ("PAGE: $user $newPages\n");

	# And now handle their historical total
	# If accFile exists: get previous count, get total, write new count
	if (-e $accFile){
		my $histPages = getUserTotal($accFile);
		my $totPages = $newPages + $histPages;
		setUserTotal($accFile,$totPages);
	}
	# Else accFile does not exist: write out new accFile
	else{
		setUserTotal($accFile,$newPages);
                # Modify read permissions so accounting is secure
                chmod (0640,"$accFile");
                system ("chown","lp","$accFile");
                system ("chgrp",$user,"$accFile"); # Note: this doesn't work unless CUPS runs this backend as root!
	}
}


sub getUserTotal{
	# Args: accounting file
	# Returns: user's historical pagecount
	my $accFile = shift;
	my $histPages;
	open (FH, "$accFile") or warn "ERROR: Cannot read $accFile: $!\n";
	while(<FH>){
		$histPages = $histPages+$_;     #extra safety in case of badly formatted file
	}
	close (FH);
	return $histPages;
}


sub setUserTotal{
	# Args: accounting file
	# Returns: nothing, writes to user's pagecount file
	my $accFile = shift;
	my $totPages = shift;
	open (OUT, ">$accFile") or warn "ERROR: Cannot write $accFile: $!\n";
	flock OUT, 2;
	print (OUT "$totPages\n");
	flock OUT, 8;
	close (OUT);
}


sub chkPrinterColor{
	# Args: printer, printerlistfile
	# Returns: 1/color printer | 0/not color printer
	my $printer = shift;
	my $accPrinterList = shift;
	open (FH, "$accPrinterList") or warn "Cannot read $accPrinterList: $!\n";
	while (<FH>){
		if (m/^$printer:=COLOR=/i){
			return 1;
		}
	}
	close FH;
	return 0;
}


sub debug{
	# Args: debugging file
	# Returns: nothing, just write out some debug info
	my $debugFile = shift;
	my $date=`/bin/date +%T%t%D`;

	open (OUT, ">>$debugFile");
        flock OUT, 2;
        seek OUT, 0, 2;
	print (OUT "\n=======================\n");

	# print time
	print (OUT "\n$date\n");

	# print args
	print (OUT "\narguments passed to backend:\n");
	my $i=0;
	foreach (@ARGV){
		print (OUT "$i:$_\t");
		$i++;
	}
	print (OUT "\n");

	# print env
	print (OUT "\nenvironment variables:\n");
	for(keys(%ENV)) {
		print (OUT "$_ = $ENV{$_}\t");
	}
	print (OUT "\n");

	print (OUT "\n=======================\n");
	flock OUT, 8;
	close (OUT);
}


sub getPrintQuota {
	# Args: accounting filedir, username
	# Returns: user's print quota
	#
	# * The user quota _always_ has first priority
	# * The group which gives the highest quota has second priority, e.g. will be used if 
	# no user quota is found.
	# * The default quota _always_ has the lowest priority.  It will only be used if no 
	# other quota is found.
	#
	# Possible return values:
	# -2   No quota files found, printing not allowed.
	# -1   Unlimited printing allowed (no quota).
	# 0    Printing explictly not allowed.
	# n>0  Quota is n pages.

        my $quotadir = shift;
        my $user = shift;
	my $acctQuota = -2; # A "no quota file found" error condition is the default.

	# First, look for a per-user quota.  This quota has top priority.
        my $quotaFile = ($quotadir,"/",$user,"_user_quota");
	if ( -e $quotaFile ) {
		$acctQuota = read_quota_file($quotaFile);
	}

	# If there is no user quota, look for a per-group quota for each group the user is in.
	# The first one found is used.  Otherwise, the default quota is still used.
        if ( ! -e $quotaFile ) {
		my ($name,$passwd,$gid,$members) = getgrent();

		while ( length($name)>0 ) {
			# This next line will allow me to use a simple pattern match.
			my $testmembers = " " . $members . " ";

			if ( $testmembers =~ /\s$user\s/ ) {
				my $testquotafile = $quotadir . "/" . $name . "_group_quota";
				my $testquota = read_quota_file($testquotafile);

				if ( $testquota == -1 ) {
					$acctQuota = $testquota;
					last;
				} elsif ( $testquota > $acctQuota ) {
					$acctQuota = $testquota;
				}
			}

			($name,$passwd,$gid,$members) = getgrent();
		}

		# If no group quota was found either, use the default quota.
		if ( length($name) == 0 ) {
	        	$quotaFile = "$quotadir/print_quota";
			$acctQuota = read_quota_file($quotaFile);
		}
	}

	return $acctQuota;
}


sub read_quota_file {
	# Arguments: file to read
	# Returns: user's current quota
	my $quotafile = shift;
	my $quota = -2; # An error condition is the default

	open (FH, "$quotafile") or return ($quota);
	while(<FH>) {
		$quota = $_;
	}
	close (FH);

	chomp($quota); # When editing files by hand, many people toss a newline in there.

	return ($quota);
}