root / tea4cups / trunk / tea4cups @ 585

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

# Tea4CUPS : Tee for CUPS
#
# (c) 2005 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$
#
#

import sys
import os
import popen2
import errno
import md5
import cStringIO
import shlex
import tempfile
import ConfigParser
from struct import unpack

version = "0.99"

class TeeError(Exception):
    """Base exception for Tea4CUPS related stuff."""
    def __init__(self, message = ""):
        self.message = message
        Exception.__init__(self, message)
    def __repr__(self):
        return self.message
    __str__ = __repr__
    
class ConfigError(TeeError) :    
    """Configuration related exceptions."""
    pass 
    
class IPPError(TeeError) :    
    """IPP related exceptions."""
    pass 
    
class Popen4ForCUPS(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):
        try :
            MAXFD = os.sysconf("SC_OPEN_MAX")
        except (AttributeError, ValueError) :    
            MAXFD = 256
        for i in range(3, MAXFD) : 
            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)
    
# Some IPP constants    
OPERATION_ATTRIBUTES_TAG = 0x01
JOB_ATTRIBUTES_TAG = 0x02
END_OF_ATTRIBUTES_TAG = 0x03
PRINTER_ATTRIBUTES_TAG = 0x04
UNSUPPORTED_ATTRIBUTES_TAG = 0x05

class IPPMessage :
    """A class for IPP message files."""
    def __init__(self, data) :
        """Initializes an IPP Message object."""
        self.data = data
        self._attributes = {}
        self.curname = None
        self.tags = [ None ] * 256      # by default all tags reserved
        
        # Delimiter tags
        self.tags[0x01] = "operation-attributes-tag"
        self.tags[0x02] = "job-attributes-tag"
        self.tags[0x03] = "end-of-attributes-tag"
        self.tags[0x04] = "printer-attributes-tag"
        self.tags[0x05] = "unsupported-attributes-tag"
        
        # out of band values
        self.tags[0x10] = "unsupported"
        self.tags[0x11] = "reserved-for-future-default"
        self.tags[0x12] = "unknown"
        self.tags[0x13] = "no-value"
        
        # integer values
        self.tags[0x20] = "generic-integer"
        self.tags[0x21] = "integer"
        self.tags[0x22] = "boolean"
        self.tags[0x23] = "enum"
        
        # octetString
        self.tags[0x30] = "octetString-with-an-unspecified-format"
        self.tags[0x31] = "dateTime"
        self.tags[0x32] = "resolution"
        self.tags[0x33] = "rangeOfInteger"
        self.tags[0x34] = "reserved-for-collection"
        self.tags[0x35] = "textWithLanguage"
        self.tags[0x36] = "nameWithLanguage"
        
        # character strings
        self.tags[0x20] = "generic-character-string"
        self.tags[0x41] = "textWithoutLanguage"
        self.tags[0x42] = "nameWithoutLanguage"
        # self.tags[0x43] = "reserved"
        self.tags[0x44] = "keyword"
        self.tags[0x45] = "uri"
        self.tags[0x46] = "uriScheme"
        self.tags[0x47] = "charset"
        self.tags[0x48] = "naturalLanguage"
        self.tags[0x49] = "mimeMediaType"
        
        # now parses the IPP message
        self.parse()
        
    def __getattr__(self, attrname) :    
        """Allows self.attributes to return the attributes names."""
        if attrname == "attributes" :
            keys = self._attributes.keys()
            keys.sort()
            return keys
        raise AttributeError, attrname
            
    def __getitem__(self, ippattrname) :    
        """Fakes a dictionnary d['key'] notation."""
        value = self._attributes.get(ippattrname)
        if value is not None :
            if len(value) == 1 :
                value = value[0]
        return value        
    get = __getitem__    
        
    def parseTag(self) :    
        """Extracts information from an IPP tag."""
        pos = self.position
        valuetag = self.tags[ord(self.data[pos])]
        # print valuetag.get("name")
        pos += 1
        posend = pos2 = pos + 2
        namelength = unpack(">H", self.data[pos:pos2])[0]
        if not namelength :
            name = self.curname
        else :    
            posend += namelength
            self.curname = name = self.data[pos2:posend]
        pos2 = posend + 2
        valuelength = unpack(">H", self.data[posend:pos2])[0]
        posend = pos2 + valuelength
        value = self.data[pos2:posend]
        oldval = self._attributes.setdefault(name, [])
        oldval.append(value)
        return posend - self.position
        
    def operation_attributes_tag(self) : 
        """Indicates that the parser enters into an operation-attributes-tag group."""
        return self.parseTag()
        
    def job_attributes_tag(self) : 
        """Indicates that the parser enters into an operation-attributes-tag group."""
        return self.parseTag()
        
    def printer_attributes_tag(self) : 
        """Indicates that the parser enters into an operation-attributes-tag group."""
        return self.parseTag()
        
    def parse(self) :
        """Parses an IPP Message.
        
           NB : Only a subset of RFC2910 is implemented.
           We are only interested in textual informations for now anyway.
        """
        self.version = "%s.%s" % (ord(self.data[0]), ord(self.data[1]))
        self.operation_id = "0x%04x" % unpack(">H", self.data[2:4])[0]
        self.request_id = "0x%08x" % unpack(">I", self.data[4:8])[0]
        self.position = 8
        try :
            tag = ord(self.data[self.position])
            while tag != END_OF_ATTRIBUTES_TAG :
                self.position += 1
                name = self.tags[tag]
                if name is not None :
                    func = getattr(self, name.replace("-", "_"), None)
                    if func is not None :
                        self.position += func()
                        if ord(self.data[self.position]) > UNSUPPORTED_ATTRIBUTES_TAG :
                            self.position -= 1
                            continue
                tag = ord(self.data[self.position])
        except IndexError :
            raise IPPError, "Unexpected end of IPP message."
            
class FakeConfig :    
    """Fakes a configuration file parser."""
    def get(self, section, option, raw=0) :
        """Fakes the retrieval of a global option."""
        raise ConfigError, "Invalid configuration file : no option %s in section [%s]" % (option, section)
        
class CupsBackend :
    """Base class for tools with no database access."""
    def __init__(self) :
        """Initializes the CUPS backend wrapper."""
        self.MyName = "Tea4CUPS"
        self.myname = "tea4cups"
        self.pid = os.getpid()
        confdir = os.environ.get("CUPS_SERVERROOT", ".") 
        self.conffile = os.path.join(confdir, "%s.conf" % self.myname)
        if os.path.isfile(self.conffile) :
            self.config = ConfigParser.ConfigParser()
            self.config.read([self.conffile])
            self.debug = self.isTrue(self.getGlobalOption("debug", ignore=1))
        else :    
            self.config = FakeConfig()
            self.debug = 1      # no config, so force debug mode !
            
    def logInfo(self, message, level="info") :        
        """Logs a message to CUPS' error_log file."""
        sys.stderr.write("%s: %s v%s (PID %i) : %s\n" % (level.upper(), self.MyName, version, os.getpid(), message))
        sys.stderr.flush()
        
    def logDebug(self, message) :    
        """Logs something to debug output if debug is enabled."""
        if self.debug :
            self.logInfo(message, level="debug")
        
    def isTrue(self, option) :        
        """Returns 1 if option is set to true, else 0."""
        if (option is not None) and (option.upper().strip() in ['Y', 'YES', '1', 'ON', 'T', 'TRUE']) :
            return 1
        else :    
            return 0
                        
    def getGlobalOption(self, option, ignore=0) :    
        """Returns an option from the global section, or raises a ConfigError if ignore is not set, else returns None."""
        try :
            return self.config.get("global", option, raw=1)
        except (ConfigParser.NoSectionError, ConfigParser.NoOptionError) :    
            if not ignore :
                raise ConfigError, "Option %s not found in section global of %s" % (option, self.conffile)
                
    def getPrintQueueOption(self, printqueuename, option, ignore=0) :    
        """Returns an option from the printer section, or the global section, or raises a ConfigError."""
        globaloption = self.getGlobalOption(option, ignore=1)
        try :
            return self.config.get(printqueuename, option, raw=1)
        except (ConfigParser.NoSectionError, ConfigParser.NoOptionError) :    
            if globaloption is not None :
                return globaloption
            elif not ignore :
                raise ConfigError, "Option %s not found in section [%s] of %s" % (option, printqueuename, self.conffile)
                
    def enumTeeBranches(self, printqueuename) :
        """Returns the list of branches for a particular section's Tee."""
        try :
            globalbranches = [ (k, v) for (k, v) in self.config.items("global") if k.startswith("tee_") ]
        except ConfigParser.NoSectionError, msg :    
            raise ConfigError, "Invalid configuration file : %s" % msg
        try :
            sectionbranches = [ (k, v) for (k, v) in self.config.items(printqueuename) if k.startswith("tee_") ]
        except ConfigParser.NoSectionError, msg :    
            self.logInfo("No section for print queue %s : %s" % (printqueuename, msg))
            sectionbranches = []
        branches = {}
        for (k, v) in globalbranches :
            value = v.strip()
            if value :
                branches[k] = value
        for (k, v) in sectionbranches :    
            value = v.strip()
            if value :
                branches[k] = value # overwrite any global option or set a new value
            else :    
                del branches[k] # empty value disables a global option
        return branches
        
    def discoverOtherBackends(self) :    
        """Discovers the other CUPS backends.
        
           Executes each existing backend in turn in device enumeration mode.
           Returns the list of available backends.
        """
        # Unfortunately this method can't output any debug information
        # to stdout or stderr, else CUPS considers that the device is
        # not available.
        available = []
        (directory, myname) = os.path.split(sys.argv[0])
        tmpdir = tempfile.gettempdir()
        lockfilename = os.path.join(tmpdir, "%s..LCK" % myname)
        if os.path.exists(lockfilename) :
            lockfile = open(lockfilename, "r")
            pid = int(lockfile.read())
            lockfile.close()
            try :
                # see if the pid contained in the lock file is still running
                os.kill(pid, 0)
            except OSError, e :    
                if e.errno != errno.EPERM :
                    # process doesn't exist anymore
                    os.remove(lockfilename)
            
        if not os.path.exists(lockfilename) :
            lockfile = open(lockfilename, "w")
            lockfile.write("%i" % self.pid)
            lockfile.close()
            allbackends = [ os.path.join(directory, b) \
                                for b in os.listdir(directory) 
                                    if os.access(os.path.join(directory, b), os.X_OK) \
                                        and (b != myname)] 
            for backend in allbackends :                            
                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(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]
                            available.append('%s %s:%s "%s+%s" "%s managed %s"' \
                                                 % (devicetype, self.myname, device, self.MyName, name, self.MyName, fullname))
            os.remove(lockfilename)
        return available
                        
    def initBackend(self) :    
        """Initializes the backend's attributes."""
        # check that the DEVICE_URI environment variable's value is 
        # prefixed with self.myname otherwise don't touch it.
        # If this is the case, we have to remove the prefix from 
        # the environment before launching the real backend 
        muststartwith = "%s:" % self.myname
        device_uri = os.environ.get("DEVICE_URI", "")
        if device_uri.startswith(muststartwith) :
            fulldevice_uri = device_uri[:]
            device_uri = fulldevice_uri[len(muststartwith):]
            for i in range(2) :
                if device_uri.startswith("/") :  
                    device_uri = device_uri[1:]
        try :
            (backend, destination) = device_uri.split(":", 1) 
        except ValueError :    
            if not device_uri :
                self.logInfo("Not attached to an existing print queue.")
                backend = ""
            else :    
                raise TeeError, "Invalid DEVICE_URI : %s\n" % device_uri
        
        self.JobId = sys.argv[1].strip()
        self.UserName = sys.argv[2].strip()
        self.Title = sys.argv[3].strip()
        self.Copies = int(sys.argv[4].strip())
        self.Options = sys.argv[5].strip()
        if len(sys.argv) == 7 :
            self.InputFile = sys.argv[6] # read job's datas from file
        else :    
            self.InputFile = None        # read job's datas from stdin
            
        self.RealBackend = backend
        self.DeviceURI = device_uri
        self.PrinterName = os.environ.get("PRINTER", "")
        self.Directory = self.getPrintQueueOption(self.PrinterName, "directory")
        self.DataFile = os.path.join(self.Directory, "%s-%s-%s-%s" % (self.myname, self.PrinterName, self.UserName, self.JobId))
        self.ClientHost = self.extractJobOriginatingHostName()
            
    def getCupsConfigDirectives(self, directives=[]) :
        """Retrieves some CUPS directives from its configuration file.
        
           Returns a mapping with lowercased directives as keys and 
           their setting as values.
        """
        dirvalues = {} 
        cupsroot = os.environ.get("CUPS_SERVERROOT", "/etc/cups")
        cupsdconf = os.path.join(cupsroot, "cupsd.conf")
        try :
            conffile = open(cupsdconf, "r")
        except IOError :    
            raise TeeError, "Unable to open %s" % cupsdconf
        else :    
            for line in conffile.readlines() :
                linecopy = line.strip().lower()
                for di in [d.lower() for d in directives] :
                    if linecopy.startswith("%s " % di) :
                        try :
                            val = line.split()[1]
                        except :    
                            pass # ignore errors, we take the last value in any case.
                        else :    
                            dirvalues[di] = val
            conffile.close()            
        return dirvalues       
            
    def extractJobOriginatingHostName(self) :        
        """Extracts the client's hostname or IP address from the CUPS message file for current job."""
        cupsdconf = self.getCupsConfigDirectives(["RequestRoot"])
        requestroot = cupsdconf.get("requestroot", "/var/spool/cups")
        if (len(self.JobId) < 5) and self.JobId.isdigit() :
            ippmessagefile = "c%05i" % int(self.JobId)
        else :    
            ippmessagefile = "c%s" % self.JobId
        ippmessagefile = os.path.join(requestroot, ippmessagefile)
        ippmessage = {}
        try :
            ippdatafile = open(ippmessagefile)
        except :    
            self.logInfo("Unable to open IPP message file %s" % ippmessagefile, "warn")
        else :    
            self.logDebug("Parsing of IPP message file %s begins." % ippmessagefile)
            try :
                ippmessage = IPPMessage(ippdatafile.read())
            except IPPError, msg :    
                self.logInfo("Error while parsing %s : %s" % (ippmessagefile, msg), "warn")
            else :    
                self.logDebug("Parsing of IPP message file %s ends." % ippmessagefile)
            ippdatafile.close()
        return ippmessage.get("job-originating-host-name")    
                
    def exportAttributes(self) :    
        """Exports our backend's attributes to the environment."""
        os.environ["DEVICE_URI"] = self.DeviceURI       # WARNING !
        os.environ["TEAPRINTERNAME"] = self.PrinterName
        os.environ["TEADIRECTORY"] = self.Directory
        os.environ["TEADATAFILE"] = self.DataFile
        os.environ["TEAJOBSIZE"] = str(self.JobSize)
        os.environ["TEAMD5SUM"] = self.JobMD5Sum
        os.environ["TEACLIENTHOST"] = self.ClientHost or ""
        os.environ["TEAJOBID"] = self.JobId
        os.environ["TEAUSERNAME"] = self.UserName
        os.environ["TEATITLE"] = self.Title
        os.environ["TEACOPIES"] = str(self.Copies)
        os.environ["TEAOPTIONS"] = self.Options
        os.environ["TEAINPUTFILE"] = self.InputFile or ""
        
    def saveDatasAndCheckSum(self) :
        """Saves the input datas into a static file."""
        self.logDebug("Duplicating data stream into %s" % self.DataFile)
        mustclose = 0
        if self.InputFile is not None :
            infile = open(self.InputFile, "rb")
            mustclose = 1
        else :    
            infile = sys.stdin
        CHUNK = 64*1024         # read 64 Kb at a time
        dummy = 0
        sizeread = 0
        checksum = md5.new()
        outfile = open(self.DataFile, "wb")    
        while 1 :
            data = infile.read(CHUNK) 
            if not data :
                break
            sizeread += len(data)    
            outfile.write(data)
            checksum.update(data)    
            if not (dummy % 32) : # Only display every 2 Mb
                self.logDebug("%s bytes saved..." % sizeread)
            dummy += 1    
        outfile.close()
        if mustclose :    
            infile.close()
        self.JobSize = sizeread    
        self.JobMD5Sum = checksum.hexdigest()
        self.logDebug("Job %s is %s bytes long." % (self.JobId, self.JobSize))
        self.logDebug("Job %s MD5 sum is %s" % (self.JobId, self.JobMD5Sum))

    def cleanUp(self) :
        """Cleans up the place."""
        if not self.isTrue(self.getPrintQueueOption(self.PrinterName, "keepfiles", ignore=1)) :
            os.remove(self.DataFile)
            
    def runBranches(self) :         
        """Launches each tee defined for the current print queue."""
        exitcode = 0
        branches = self.enumTeeBranches(self.PrinterName)
        if self.isTrue(self.getPrintQueueOption(self.PrinterName, "serialize", ignore=1)) :
            self.logDebug("Serialized Tees")
            for (branch, command) in branches.items() :
                self.logDebug("Launching %s : %s" % (branch, command))
                retcode = os.system(command)
                self.logDebug("Exit code for tee %s on printer %s is %s" % (branch, self.PrinterName, retcode))
                if os.WIFEXITED(retcode) :
                    retcode = os.WEXITSTATUS(retcode)
                if retcode :    
                    self.logInfo("Tee %s on printer %s didn't exit successfully." % (branch, self.PrinterName), "error")
                    exitcode = 1
        else :        
            self.logDebug("Forked Tees")
            pids = {}
            for (branch, command) in branches.items() :
                pid = os.fork()
                if pid :
                    pids[branch] = pid
                else :    
                    self.logDebug("Launching %s : %s" % (branch, command))
                    retcode = os.system(command)
                    if os.WIFEXITED(retcode) :
                        retcode = os.WEXITSTATUS(retcode)
                    else :    
                        retcode = -1
                    sys.exit(retcode)
            for (branch, pid) in pids.items() :
                (childpid, retcode) = os.waitpid(pid, 0)
                self.logDebug("Exit code for tee %s (PID %s) on printer %s is %s" % (branch, childpid, self.PrinterName, retcode))
                if os.WIFEXITED(retcode) :
                    retcode = os.WEXITSTATUS(retcode)
                if retcode :    
                    self.logInfo("Tee %s (PID %s) on printer %s didn't exit successfully." % (branch, childpid, self.PrinterName), "error")
                    exitcode = 1
        return exitcode
        
if __name__ == "__main__" :    
    # This is a CUPS backend, we should act and die like a CUPS backend
    wrapper = CupsBackend()
    if len(sys.argv) == 1 :
        print "\n".join(wrapper.discoverOtherBackends())
        sys.exit(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])
        sys.exit(1)
    else :    
        try :
            wrapper.initBackend()
            wrapper.saveDatasAndCheckSum()
            wrapper.exportAttributes()
            retcode = wrapper.runBranches()
            wrapper.cleanUp()
        except SystemExit, e :    
            retcode = e.code
        except :    
            import traceback
            lines = []
            for line in traceback.format_exception(*sys.exc_info()) :
                lines.extend([l for l in line.split("\n") if l])
            msg = "ERROR: ".join(["%s (PID %s) : %s\n" % (wrapper.MyName, wrapper.pid, l) for l in (["ERROR: Tea4CUPS v%s" % version] + lines)])
            sys.stderr.write(msg)
            sys.stderr.flush()
            retcode = 1
        sys.exit(retcode)