#!/usr/bin/perl

##################################################################
# A very simple quota enforcement script for use with lprng.
# This script was written for at-home use, to restrict the
# number of pages that my young son can print each day.
#
# Written 2010-10-28 by Lester Hightower
##################################################################

use strict;
use Getopt::Std;
use FileHandle;
use Data::Dumper;
use Config::Tiny;
use POSIX qw(strftime);

my $VERSION = 0.3;

our %EXE = (
	'pclcount' => '/usr/lib/lprng/filters/pclcount',
	'file' => '/usr/bin/file',
);
my($JFAIL, $JABORT, $JREMOVE, $JHOLD) = ( 32, 33, 34, 37);
my %opt;
getopts( 'A:B:C:D:E:F:G:H:I:J:K:L:M:N:O:P:Q:R:T:S:U:V:W:X:Y:Z:'
	. 'a:b:cd:e:f:g:h:i:j:k:l:m:n:o:p:q:r:t:s:u:v:w:x:y:z:', \%opt );

# $DH is the "debug handle" / our log-file. $DH can be set to
# STDERR for logging to go to /var/spool/lpd/<printer>/log.
our $DH=new FileHandle;
open($DH,">> $opt{d}/lpquota.log");

# The meat of the entire program is these three functions
my $config = get_quota_config(\%opt);
my $job_hfA = get_job_hfA(\%opt);
my $allow_this_job = do_accounting(\%opt,$config,$job_hfA);

if (0) { # Used for debugging only
  foreach $a (qw(SPOOL_DIR CONTROL_DIR PRINTCAP_ENTRY CONTROL)) {
    print $DH "$a=$ENV{$a}\n";
  }
  print $DH "\n\n" . join(", ", @ARGV);
  print $DH "\n\n" . Dumper(\%opt);
  print $DH "\n\n" . Dumper($job_hfA);
}

close $DH;

# Exit-code based on allowing this job to proceed or not
if ($allow_this_job) { exit 0; }
exit $JREMOVE;

########################################################################
# FUNCTIONS ############################################################
########################################################################

sub do_accounting($$$) {
  my $opt = shift @_;
  my $config = shift @_;
  my $job_hfA = shift @_;

  my $user = $job_hfA->{'P'};
  if (! defined($config->{$user})) {
    return 1; # This user is not restricted on this printer
  }

  my $data_store_file = $config->{$user}->{data_store};
  my $data_store;
  if (-e $data_store_file) {
    $data_store = Config::Tiny->read($data_store_file);
  } else {
    $data_store = new Config::Tiny;
  }

  my $today=strftime("%Y%m%d", localtime);
  my $spent_pages=0;
  if (defined($data_store->{_}->{date}) &&
	$data_store->{_}->{date} eq $today) {
    $spent_pages = $data_store->{_}->{pages};
  } else {
    $data_store->{_}->{date} = $today;
    $spent_pages=0;
  }

  my $max_daily_pages = $config->{$user}->{max_daily_pages};
  print $DH "User $user is quota-restricted to $max_daily_pages pages/day!!\n";
  # Quota exceeded
  if (($spent_pages + $job_hfA->{pages}) > $max_daily_pages) {
    print $DH "User $user has exhausted the quota for $today!!\n";
    print $DH " - This job had $job_hfA->{pages} pages\n";
    print $DH " - Spent pages prior to this job was $spent_pages\n";
    return 0;
  } else {
    # Record new spent pages
    $data_store->{_}->{pages} = $spent_pages + $job_hfA->{pages};
    $data_store->write($data_store_file);
    print $DH " - This job had $job_hfA->{pages} pages\n";
    print $DH " - New spent pages is " . $data_store->{_}->{pages} . "\n";
    return 1; # Allow this job to go through for this restricted user
  }
}

sub keyval_file_to_hash($) {
  my $filename=shift @_;
  my $fh=new FileHandle;
  my %h=();
  if (open($fh, "< $filename")) {
    foreach my $line (<$fh>) {
      chomp $line;
      my ($key,$val) = split('=', $line, 2);
      $h{$key} = $val;
    }
    close $fh;
  }
  return \%h;
}

sub get_quota_config($) {
  my $opt = shift @_;
  my $printer = $opt->{'Q'};
  my $quota_config_file = "/etc/lpquota_" . $printer . ".conf";
  my $config = {};
  if (-r $quota_config_file) {
    $config = Config::Tiny->read($quota_config_file);
  }
  return $config;
}

sub get_job_hfA($) {
  my $opt = shift @_;
  my $job_number = $opt->{'j'};
  my $spool_dir = $opt->{'d'};
  my $job_control = $spool_dir . '/hfA' . $job_number;
  my $job_hfA = keyval_file_to_hash($job_control);

  # Add the page count data
  my $datafile = $spool_dir . '/' . $job_hfA->{'datafiles'};
  $datafile =~ s/\s+$//;
  my $type=`'$EXE{file}' -b '$datafile'`;
  my $pages = 0;
  if ($type =~ m/^HP Printer Job Language/i) {
    $pages=`'$EXE{pclcount}' '$datafile'`;
  }
  # TODO: Add postscript support. Should be as easy as 'grep -c showpage'
  $job_hfA->{'pages'} = int($pages);

  return $job_hfA;
}


########################################################################
# POD ##################################################################
########################################################################

=head1 NAME

lpquota.pl - A very simple quota enforcement script for use with lprng.


=head1 DESCRIPTION

This script was written for at-home use, to restrict the number of pages that
my young son can print each day.


=head1 PREREQUISITES

This script requires the non-core module C<Config::Tiny>.
 - "sudo apt-get install libconfig-tiny-perl" on Ubuntu.

For PCL printers, this script depends on pclcount:
 - http://www.fea.unicamp.br/pclcount/


=head1 CONFIGURATION

 1) Properly set the %EXE values in the top of the script.
 2) Add a line like ":as=|/usr/lib/lprng/filters/lpquota.pl:"
    to your /etc/printcap for the desired printer(s).
 3) Create a file /etc/lpquota_<printer>.conf that looks like:
    # cat /etc/lpquota_lp.conf 
    ; Limit Andrew's printing (Andrew is his WinXP username)
    ; tail /var/spool/lpd/lp0/acct and look for "-nUsername"
    [Andrew]
    max_daily_pages=5
    data_store=/var/cache/lpquota/Andrew_lp
 4) If you follow my example in step 3, then make the directory
    /var/cache/lpquota and set proper permissions. On Ubuntu it
    will need to be owned by daemon.lp and writeable by them.

=pod OSNAMES

Unix-like (written and tested on Ubuntu Linux 10.04 LTS).

=pod SCRIPT CATEGORIES

UNIX/System_administration

=cut