root / pykota / trunk / bin / cupspykota @ 1478

#! /usr/bin/env python
# -*- coding: ISO-8859-15 -*-

# CUPSPyKota accounting backend
#
# PyKota - Print Quotas for CUPS and LPRng
#
# (c) 2003-2004 Jerome Alet <alet@librelogiciel.com>
# 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.
#
# $Id$
#
# $Log$
# Revision 1.43  2004/05/17 11:46:05  jalet
# First try at cupspykota's main loop rewrite
#
# Revision 1.42  2004/05/10 11:22:28  jalet
# Typo
#
# Revision 1.41  2004/05/10 10:07:30  jalet
# Catches OSError while reading
#
# Revision 1.40  2004/05/10 09:29:48  jalet
# Should be more robust if we receive a SIGTERM during an I/O operation
#
# Revision 1.39  2004/05/07 14:44:53  jalet
# Fix for file handles unregistered twice from the polling object
#
# Revision 1.38  2004/04/09 22:24:46  jalet
# Began work on correct handling of child processes when jobs are cancelled by
# the user. Especially important when an external requester is running for a
# long time.
#
# Revision 1.37  2004/03/18 19:11:25  jalet
# Fix for raw jobs in cupspykota
#
# Revision 1.36  2004/03/18 14:03:18  jalet
# Added fsync() calls
#
# Revision 1.35  2004/03/16 12:05:01  jalet
# Small fix for new waitprinter.sh : when job was denied, would wait forever
# for printer being in printing mode.
#
# Revision 1.34  2004/03/15 10:47:56  jalet
# This time the traceback formatting should be correct !
#
# Revision 1.33  2004/03/05 12:46:07  jalet
# Improve tracebacks
#
# Revision 1.32  2004/03/05 12:31:35  jalet
# Now should output full traceback when crashing
#
# Revision 1.31  2004/03/01 14:35:56  jalet
# PYKOTAPHASE wasn't set soon enough at the start of the job
#
# Revision 1.30  2004/03/01 14:34:15  jalet
# PYKOTAPHASE wasn't set at the right time at the end of data transmission
# to underlying layer (real backend)
#
# Revision 1.29  2004/03/01 11:23:25  jalet
# Pre and Post hooks to external commands are available in the cupspykota
# backend. Forthe pykota filter they will be implemented real soon now.
#
# Revision 1.28  2004/02/26 14:18:07  jalet
# Should fix the remaining bugs wrt printers groups and users groups.
#
# Revision 1.27  2004/02/04 23:41:27  jalet
# Should fix the incorrect "backend died abnormally" problem.
#
# Revision 1.26  2004/01/30 16:35:03  jalet
# Fixes stupid software accounting bug in CUPS backend
#
# Revision 1.25  2004/01/16 17:51:46  jalet
# Fuck Fuck Fuck !!!
#
# Revision 1.24  2004/01/14 15:52:01  jalet
# Small fix for job cancelling code.
#
# Revision 1.23  2004/01/13 10:48:28  jalet
# Small streams polling loop modification.
#
# Revision 1.22  2004/01/12 22:43:40  jalet
# New formula to compute a job's price
#
# Revision 1.21  2004/01/12 18:17:36  jalet
# Denied jobs weren't stored into the history anymore, this is now fixed.
#
# Revision 1.20  2004/01/11 23:22:42  jalet
# Major code refactoring, it's way cleaner, and now allows automated addition
# of printers on first print.
#
# Revision 1.19  2004/01/08 14:10:32  jalet
# Copyright year changed.
#
# Revision 1.18  2004/01/07 16:16:32  jalet
# Better debugging information
#
# Revision 1.17  2003/12/27 16:49:25  uid67467
# Should be ok now.
#
# Revision 1.17  2003/12/06 08:54:29  jalet
# Code simplifications.
# Added many debugging messages.
#
# Revision 1.16  2003/11/26 20:43:29  jalet
# Inadvertantly introduced a bug, which is fixed.
#
# Revision 1.15  2003/11/26 19:17:35  jalet
# Printing on a printer not present in the Quota Storage now results
# in the job being stopped or cancelled depending on the system.
#
# Revision 1.14  2003/11/25 13:25:45  jalet
# Boolean problem with old Python, replaced with 0
#
# Revision 1.13  2003/11/23 19:01:35  jalet
# Job price added to history
#
# Revision 1.12  2003/11/21 14:28:43  jalet
# More complete job history.
#
# Revision 1.11  2003/11/19 23:19:35  jalet
# Code refactoring work.
# Explicit redirection to /dev/null has to be set in external policy now, just
# like in external mailto.
#
# Revision 1.10  2003/11/18 17:54:24  jalet
# SIGTERMs are now transmitted to original backends.
#
# Revision 1.9  2003/11/18 14:11:07  jalet
# Small fix for bizarre urls
#
# Revision 1.8  2003/11/15 14:26:44  jalet
# General improvements to the documentation.
# Email address changed in sample configuration file, because
# I receive low quota messages almost every day...
#
# Revision 1.7  2003/11/14 22:05:12  jalet
# New CUPS backend fully functionnal.
# Old CUPS configuration method is now officially deprecated.
#
# Revision 1.6  2003/11/14 20:13:11  jalet
# We exit the loop too soon.
#
# Revision 1.5  2003/11/14 18:31:27  jalet
# Not perfect, but seems to work with the poll() loop.
#
# Revision 1.4  2003/11/14 17:04:15  jalet
# More (untested) work on the CUPS backend.
#
# Revision 1.3  2003/11/12 23:27:44  jalet
# More work on new backend. This commit may be unstable.
#
# Revision 1.2  2003/11/12 09:33:34  jalet
# New CUPS backend supports device enumeration
#
# Revision 1.1  2003/11/08 16:05:31  jalet
# CUPS backend added for people to experiment.
#
#
#

import sys
import os
import fcntl
import popen2
import cStringIO
import shlex
import select
import signal
import time

from pykota.tool import PyKotaFilterOrBackend, PyKotaToolError
from pykota.config import PyKotaConfigError
from pykota.storage import PyKotaStorageError
from pykota.accounter import PyKotaAccounterError
from pykota.requester import PyKotaRequesterError
    
class PyKotaPopen4(popen2.Popen4) :
    """Our own class to execute real backends.
    
       Their first argument is different from their path so using
       native popen2.Popen3 would not be feasible.
    """
    def __init__(self, cmd, bufsize=-1, arg0=None) :
        self.arg0 = arg0
        popen2.Popen4.__init__(self, cmd, bufsize)
        
    def _run_child(self, cmd):
        for i in range(3, 256): # TODO : MAXFD in original popen2 module
            try:
                os.close(i)
            except OSError:
                pass
        try:
            os.execvpe(cmd[0], [self.arg0 or cmd[0]] + cmd[1:], os.environ)
        finally:
            os._exit(1)
    
class PyKotaBackend(PyKotaFilterOrBackend) :        
    """A class for the pykota backend."""
    def acceptJob(self) :        
        """Returns the appropriate exit code to tell CUPS all is OK."""
        return 0
            
    def removeJob(self) :            
        """Returns the appropriate exit code to let CUPS think all is OK.
        
           Returning 0 (success) prevents CUPS from stopping the print queue.
        """   
        return 0
        
    def doWork(self, policy, printer, user, userpquota) :    
        """Most of the work is done here."""
        # Two different values possible for policy here :
        # ALLOW means : Either printer, user or user print quota doesn't exist,
        #               but the job should be allowed anyway.
        # OK means : Both printer, user and user print quota exist, job should
        #            be allowed if current user is allowed to print on this printer
        if policy == "OK" :
            # exports user information with initial values
            self.exportUserInfo(userpquota)
            
            # enters first phase
            os.putenv("PYKOTAPHASE", "BEFORE")
            
            # checks the user's quota
            action = self.warnUserPQuota(userpquota)
            
            # exports some new environment variables
            os.putenv("PYKOTAACTION", action)
            
            # launches the pre hook
            self.prehook(userpquota)
            
            self.logdebug("Job accounting begins.")
            self.accounter.beginJob(userpquota)
        else :    
            action = "ALLOW"
            os.putenv("PYKOTAACTION", action)
            
        # pass the job's data to the real backend    
        if action in ["ALLOW", "WARN"] :
            if self.gotSigTerm :
                retcode = self.removeJob()
            else :    
                retcode = self.handleData()        
        else :        
            retcode = self.removeJob()
        
        if policy == "OK" :        
            # indicate phase change
            os.putenv("PYKOTAPHASE", "AFTER")
            
            # stops accounting. 
            self.accounter.endJob(userpquota)
            self.logdebug("Job accounting ends.")
                
            # retrieve the job size    
            if action == "DENY" :
                jobsize = 0
                self.logdebug("Job size forced to 0 because printing is denied.")
            else :    
                jobsize = self.accounter.getJobSize()
            self.logdebug("Job size : %i" % jobsize)
            
            # update the quota for the current user on this printer 
            self.logdebug("Updating user %s's quota on printer %s" % (user.Name, printer.Name))
            jobprice = userpquota.increasePagesUsage(jobsize)
            
            # adds the current job to history    
            printer.addJobToHistory(self.jobid, user, self.accounter.getLastPageCounter(), action, jobsize, jobprice, self.preserveinputfile, self.title, self.copies, self.options)
            self.logdebug("Job added to history.")
            
            # exports some new environment variables
            os.putenv("PYKOTAJOBSIZE", str(jobsize))
            os.putenv("PYKOTAJOBPRICE", str(jobprice))
            
            # then re-export user information with new values
            self.exportUserInfo(userpquota)
            
            # Launches the post hook
            self.posthook(userpquota)
            
        return retcode    
               
    def setNonBlocking(self, fno) :
        """Sets a file handle to be non-blocking."""
        flags = fcntl.fcntl(fno, fcntl.F_GETFL, 0)
        fcntl.fcntl(fno, fcntl.F_SETFL, flags | os.O_NONBLOCK)

    def unregisterFileNo(self, pollobj, fileno) :                
        """Removes a file handle from the polling object."""
        try :
            pollobj.unregister(fileno)
        except KeyError :    
            self.logger.log_message(_("File number %s unregistered twice from polling object, ignored.") % fileno, "warn")
        else :    
            self.logdebug("File number %s unregistered from polling object." % fileno)
            
    def formatFileEvent(self, fd, mask, ins, outs) :        
        """Formats file debug info."""
        try :
            name = ins.get(fd, outs.get(fd))["name"]
        except KeyError :    
            self.logdebug("File %s not found in %s or %s" % (fd, repr(ins), repr(outs)))
        else :    
            maskval = []
            if mask & select.POLLIN :
                maskval.append("POLLIN")
            if mask & select.POLLOUT :
                maskval.append("POLLOUT")
            if mask & select.POLLPRI :
                maskval.append("POLLPRI")
            if mask & select.POLLERR :
                maskval.append("POLLERR")
            if mask & select.POLLHUP :
                maskval.append("POLLHUP")
            if mask & select.POLLNVAL :
                maskval.append("POLLNVAL")
            return "%s (%s)" % (name, " | ".join(maskval))
        
    def handleData(self) :                    
        """Pass the job's data to the real backend."""
        # Find the real backend pathname    
        realbackend = os.path.join(os.path.split(sys.argv[0])[0], self.originalbackend)
        
        # And launch it
        self.logdebug("Starting real backend %s with args %s" % (realbackend, " ".join(['"%s"' % a for a in ([os.environ["DEVICE_URI"]] + sys.argv[1:])])))
        subprocess = PyKotaPopen4([realbackend] + sys.argv[1:], bufsize=0, arg0=os.environ["DEVICE_URI"])
        
        # Save file descriptors, we will need them later.
        stderrfno = sys.stderr.fileno()
        fromcfno = subprocess.fromchild.fileno()
        self.setNonBlocking(fromcfno)
        
        # We will have to be careful when dealing with I/O 
        # So we use a poll object to know when to read or write
        pollster = select.poll()
        pollster.register(fromcfno, select.POLLIN | select.POLLPRI)
        
        instreams = { \
                      fromcfno : { "file" : subprocess.fromchild, "out" : stderrfno, "done" : 0, "name" : "real backend's stdout+stderr" },\
                    }
                      
        outstreams = { \
                       stderrfno : { "file" : sys.stderr, "done" : 0, "in" : fromcfno, "name" : "stderr" }, \
                     }
                       
        if self.preserveinputfile is None :
            # this is not a real file, we read the job's data 
            # from stdin and send it on our stdout
            tocfno = subprocess.tochild.fileno()
            stdinfno = sys.stdin.fileno()
            self.setNonBlocking(stdinfno)
            pollster.register(stdinfno, select.POLLIN | select.POLLPRI)
            instreams.update({ stdinfno : { "file": sys.stdin, "out" : tocfno, "done" : 0, "name" : "stdin" }})
            outstreams.update({ tocfno : { "file" : subprocess.tochild, "done" : 0, "in" : stdinfno, "name" : "real backend's stdin" }})
        else :    
            # job's data is in a file, no need to pass the data
            # to the real backend
            self.logdebug("Job's data is in file %s" % self.preserveinputfile)
            
        killed = 0
        status = -1
        self.logdebug("Entering streams polling loop...")
        while status == -1 :
            # Catches IOErrors caused by interrupted system calls
            try :
                # First check if original backend is still alive
                status = subprocess.poll()
                
                # Now if we got SIGTERM, we have 
                # to kill -TERM the original backend
                if self.gotSigTerm and not killed :
                    os.kill(subprocess.pid, signal.SIGTERM)
                    self.logger.log_message(_("SIGTERM was sent to real backend %s (pid: %s)") % (realbackend, subprocess.pid), "info")
                    killed = 1
                
                # In any case, deal with any remaining I/O
                availablefds = pollster.poll(5000)
                if not availablefds :
                    self.logdebug("Nothing to do, sleeping a bit...")
                    time.sleep(0.01) # nothing to do, give time to CPU
                else :    
                    for (fd, mask) in availablefds :
                        # self.logdebug(self.formatFileEvent(fd, mask, instreams, outstreams))
                        try :
                            if mask & (select.POLLIN | select.POLLPRI) :     
                                # We have something to read
                                try :
                                    fobj = instreams[fd]
                                except KeyError :    
                                    self.logdebug("READ : %s" % self.formatFileEvent(fd, mask, instreams, outstreams))
                                else :    
                                    data = fobj["file"].read()
                                    if not data :
                                        self.logdebug("No more data to read on %s (read returned nothing)" % fobj["name"])
                                        if not fobj["done"] :
                                            self.unregisterFileNo(pollster, fd)
                                            fobj["done"] = 1
                                    else :    
                                        # self.logdebug("%s -- DATA[%i] <= : %s ..." % (self.formatFileEvent(fd, mask, instreams, outstreams), len(data), data[:50]))
                                        fout = outstreams[fobj["out"]]["file"]
                                        fout.write(data)
                                        fout.flush()
                                            
                            if mask & (select.POLLHUP | select.POLLERR) :
                                # Some pipe has no more datas so we don't
                                # want to continue to poll this file
                                toclose = None
                                try :
                                    fobj = instreams[fd]
                                    if fobj["name"] == "stdin" :
                                        toclose = outstreams[fobj["out"]]
                                    self.logdebug("No more data to read from %s (POLLUP or POLLERR received)" % fobj["name"])
                                except KeyError :    
                                    fobj = outstreams[fd]
                                    if fobj["name"] == "stderr" :
                                        toclose = instreams[fobj["in"]]
                                    self.logdebug("No more data to write to %s (POLLUP or POLLERR received)" % fobj["name"])
                                    
                                if not fobj["done"] :
                                    self.unregisterFileNo(pollster, fd)
                                    fobj["done"] = 1
                                    if toclose is not None :
                                        self.logdebug("Closing %s" % toclose["name"])
                                        toclose["file"].close()
                                        
                            if mask & select.POLLNVAL :            
                                self.logdebug("CLOSED : %s" % self.formatFileEvent(fd, mask, instreams, outstreams))
                                
                        except IOError, msg :                
                            self.logdebug("IOError : %s -- %s" % (msg, self.formatFileEvent(fd, mask, instreams, outstreams)))
                            time.sleep(0.01) # give some time to the CPU
            except IOError, msg :
                self.logdebug("IOError : %s" % msg)
                time.sleep(0.01) # give some time to the CPU
                
        self.logdebug("Exiting streams polling loop...")
        status = subprocess.wait()      # just in case
        if os.WIFEXITED(status) :
            retcode = os.WEXITSTATUS(status)
        elif not killed :    
            self.logger.log_message(_("CUPS backend %s died abnormally.") % realbackend, "error")
            retcode = -1
        else :    
            retcode = self.removeJob()
        self.logdebug("Real backend exited with status %s" % status)    
        return retcode    
    
if __name__ == "__main__" :    
    # This is a CUPS backend, we should act and die like a CUPS backend
    if len(sys.argv) == 1 :
        # we will execute each existing backend in device enumeration mode
        # and generate their PyKota accounting counterpart
        (directory, myname) = os.path.split(sys.argv[0])
        for backend in [os.path.join(directory, b) for b in os.listdir(directory) if os.path.isfile(os.path.join(directory, b)) and (b != myname)] :
            answer = os.popen(backend, "r")
            try :
                devices = [line.strip() for line in answer.readlines()]
            except :    
                devices = []
            status = answer.close()
            if status is None :
                for d in devices :
                    # each line is of the form : 'xxxx xxxx "xxxx xxx" "xxxx xxx"'
                    # so we have to decompose it carefully
                    fdevice = cStringIO.StringIO("%s" % d)
                    tokenizer = shlex.shlex(fdevice)
                    tokenizer.wordchars = tokenizer.wordchars + r".:,?!~/\_$*-+={}[]()#"
                    arguments = []
                    while 1 :
                        token = tokenizer.get_token()
                        if token :
                            arguments.append(token)
                        else :
                            break
                    fdevice.close()
                    try :
                        (devicetype, device, name, fullname) = arguments
                    except ValueError :    
                        pass    # ignore this 'bizarre' device
                    else :    
                        if name.startswith('"') and name.endswith('"') :
                            name = name[1:-1]
                        if fullname.startswith('"') and fullname.endswith('"') :
                            fullname = fullname[1:-1]
                        print '%s cupspykota:%s "PyKota+%s" "PyKota managed %s"' % (devicetype, device, name, fullname)
        retcode = 0
    elif len(sys.argv) not in (6, 7) :    
        sys.stderr.write("ERROR: %s job-id user title copies options [file]\n" % sys.argv[0])
        retcode = 1
    else :    
        try :
            # Initializes the backend
            kotabackend = PyKotaBackend()    
            retcode = kotabackend.mainWork()
        except (PyKotaToolError, PyKotaConfigError, PyKotaStorageError, PyKotaAccounterError, PyKotaRequesterError, AttributeError, KeyError, IndexError, ValueError, TypeError, IOError), msg :
            import traceback
            mm = [((f.endswith('\n') and f) or (f + '\n')) for f in traceback.format_exception(*sys.exc_info())]
            sys.stderr.write("ERROR : cupspykota backend failed (%s)\n%s" % (msg, "ERROR : ".join(mm)))
            sys.stderr.flush()
            retcode = 1
        
        try :
            kotabackend.storage.close()
        except (TypeError, NameError, AttributeError) :    
            pass
        
    sys.exit(retcode)