root / pykota / trunk / bin / edpykota @ 2657

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

# PyKota Print Quota Editor 
#
# PyKota - Print Quotas for CUPS and LPRng
#
# (c) 2003, 2004, 2005, 2006 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., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# $Id$
#
#

import sys
import os
import pwd
import grp
from pykota.tool import PyKotaTool, PyKotaToolError, PyKotaCommandLineError, crashed, N_
from pykota.config import PyKotaConfigError
from pykota.storage import PyKotaStorageError

__doc__ = N_("""edpykota v%(__version__)s (c) %(__years__)s %(__author__)s

A Print Quota editor for PyKota.

command line usage :

  edpykota [options] user1 user2 ... userN
  
  edpykota [options] group1 group2 ... groupN

options :

  -v | --version       Prints edpykota's version number then exits.
  -h | --help          Prints this message then exits.
  
  -a | --add           Adds users and/or printers if they don't 
                       exist on the Quota Storage Server.
                       
  -d | --delete        Deletes users/groups from the quota storage.
                       Printers are never deleted.
                       
  -c | --charge p[,j]  Sets the price per page and per job to charge
                       for a particular printer. Job price is optional.
                       If both are to be set, separate them with a comma.
                       Floating point values are allowed.
                       
  -o | --overcharge f  Sets the overcharging factor applied to the user 
                       when computing the cost of a print job. Positive or 
                       negative floating point values are allowed,
                       this allows you to do some really creative
                       things like giving money to an user whenever
                       he prints. The number of pages in a print job
                       is not modified by this coefficient, only the
                       cost of the job for a particular user.
                       Only users have a coefficient.
  
  -i | --ingroups g1[,g2...]  Puts the users into each of the groups
                              listed, separated by commas. The groups
                              must already exist in the Quota Storage.
  
  -u | --users         Edit users print quotas, this is the default.
  
  -P | --printer p     Edit quotas on printer p only. Actually p can
                       use wildcards characters to select only
                       some printers. The default value is *, meaning
                       all printers. 
                       You can specify several names or wildcards, 
                       by separating them with commas.
  
  -G | --pgroups pg1[,pg2...] Adds the printer(s) to the printer groups
                       pg1, pg2, etc... which must already exist.
                       A printer group is just like a normal printer,
                       only that it is usually unknown from the printing
                       system. Create printer groups exactly the same
                       way that you create printers, then add other 
                       printers to them with this option.
                       Accounting is done on a printer and on all
                       the printer groups it belongs to, quota checking
                       is done on a printer and on all the printer groups
                       it belongs to.
  
  -g | --groups        Edit users groups print quotas instead of users.
                          
  -p | --prototype u|g Uses user u or group g as a prototype to set
                       print quotas
                       
  -n | --noquota       Sets both soft and hard limits to None for a
                       particular print quota entry.
                       This is NOT the same as --limitby noquota
                       which acts on ALL print quota entries for a
                       particular user.
  
  -r | --reset         Resets the actual page counter for the user
                       or group to zero on the specified printers. 
                       The life time page counter is kept unchanged.
                       
  -R | --hardreset     Resets the actual and life time page counters
                       for the user or group to zero on the specified 
                       printers. This is a shortcut for '--used 0'.
                       
  -l | --limitby l     Choose if the user/group is limited in printing                     
                       by its account balance or by its page quota.
                       The default value is 'quota'. Allowed values
                       are 'quota' 'balance' 'noquota' 'noprint' 
                       and 'nochange' :
                       
                         - quota : limit by number of pages per printer.
                         - balance : limit by number of credits in account.
                         - noquota : no limit, accounting still done.
                         - nochange : no limit, accounting not done. 
                         - noprint : printing is denied. 
                       NB : nochange and noprint are not supported for groups.
                       
  -b | --balance b     Sets the user's account balance to b.                     
                       Account balance may be increase or decreased
                       if b is prefixed with + or -.
                       WARNING : when decreasing account balance,
                       the total paid so far by the user is decreased
                       too.
                       Groups don't have a real balance, but the
                       sum of their users' account balance.
                       
  -C | --comment txt   Defines some informational text to be associated
                       with a change to an user's account balance.
                       Only meaningful if -b | --balance is also used.
                       
  -S | --softlimit sl  Sets the quota soft limit to sl pages.                       
  
  -H | --hardlimit hl  Sets the quota hard limit to hl pages.
  
  -I | --increase v    Increase both Soft and Hard limits by the value
                       of v. You can prefix v with + or -, if no sign is
                       used, + is assumed.

  -U | --used usage    Sets the pagecounters for the user to usage pages;
                       useful for migrating users from a different system
                       where they have already used some pages. Actual
                       and Life Time page counters may be increased or decreased
                       if usage is prefixed with + or -.
                       WARNING : BOTH page counters are modified in all cases,
                       so be careful.
                       NB : if 'usage' equals '0', then the action taken is
                       the same as if --hardreset was used.

  user1 through userN and group1 through groupN can use wildcards
  if the --add option is not set.
  
examples :                              

  $ edpykota --add -p jerome john paul george ringo/ringo@example.com
  
  This will add users john, paul, george and ringo to the quota
  database, and set their print quotas to the same values than user 
  jerome. User jerome must already exist.
  User ringo's email address will also be set to 'ringo@example.com'
  
  $ edpykota --printer lp -S 50 -H 60 jerome
  
  This will set jerome's print quota on the lp printer to a soft limit
  of 50 pages, and a hard limit of 60 pages. If either user jerome or
  printer lp doesn't exist on the Quota Storage Server then nothing is done.

  $ edpykota --add --printer lp --ingroups coders,it -S 50 -H 60 jerome
  
  Same as above, but if either user jerome or printer lp doesn't exist 
  on the Quota Storage Server they are automatically added. Also
  user jerome is put into the groups "coders" and "it" which must
  already exist in the Quota Storage.
            
  $ edpykota -g -S 500 -H 550 financial support            
  
  This will set print quota soft limit to 500 pages and hard limit
  to 550 pages for groups financial and support on all printers.
  
  $ edpykota --reset jerome "jo*"
  
  This will reset jerome's page counter to zero on all printers, as
  well as every user whose name begins with 'jo'.
  Their life time page counter on each printer will be kept unchanged.
  You can also reset the life time page counters by using the
  --hardreset | -R command line option.
  
  $ edpykota --printer hpcolor --noquota jerome
  
  This will tell PyKota to not limit jerome when printing on the 
  hpcolor printer. All his jobs will be allowed on this printer, but 
  accounting of the pages he prints will still be kept.
  Print Quotas for jerome on other printers are unchanged.
  
  $ edpykota --limitby balance jerome
  
  This will tell PyKota to limit jerome by his account's balance
  when printing.
  
  $ edpykota --balance +10.0 jerome
  
  This will increase jerome's account balance by 10.0 (in your
  own currency). You can decrease the account balance with a
  dash prefix, and set it to a fixed amount with no prefix.
  
  $ edpykota --delete jerome rachel
  
  This will completely delete jerome and rachel from the Quota Storage
  database. All their quotas and jobs will be deleted too.
  
  $ edpykota --printer lp --charge 0.1
  
  This will set the page price for printer lp to 0.1. Job price
  will not be changed.
  
  $ edpykota --printer hplj1,hplj2 --pgroups Laser,HP
  
  This will put printers hplj1 and hplj2 in printers groups Laser and HP.
  When printing either on hplj1 or hplj2, print quota will also be 
  checked and accounted for on virtual printers Laser and HP.
  
  $ edpykota --overcharge 2.5 poorstudent
  
  This will overcharge the poorstudent user by a factor of 2.5.
  
  $ edpykota --overcharge -1 jerome
  
  User jerome will actually earn money whenever he prints.
  
  $ edpykota --overcharge 0 boss
  
  User boss can print at will, it won't cost him anything because the
  cost of each print job will be multiplied by zero before charging
  his account.
""") 
        
class EdPyKota(PyKotaTool) :        
    """A class for edpykota."""
    def main(self, names, options) :
        """Edit user or group quotas."""
        if not self.config.isAdmin :
            raise PyKotaCommandLineError, "%s : %s" % (pwd.getpwuid(os.geteuid())[0], _("You're not allowed to use this command."))
        
        suffix = (options["groups"] and "Group") or "User"        
        
        if options["delete"] :    
            if not names :
                raise PyKotaCommandLineError, _("You have to pass user or group names on the command line")
            self.display("Processing...\n")
            todelete = getattr(self.storage, "getMatching%ss" % suffix)(",".join(names))
            nbtotal = len(todelete)
            for i in range(nbtotal) :
                todelete[i].delete()
                percent = 100.0 * float(i) / float(nbtotal)
                self.display("\r%.02f%%" % percent)
            self.display("\r100.00%%\r        ")
            self.display("\r%s\n" % _("Done."))
        else :
            softlimit = hardlimit = None
            
            limitby = options["limitby"]
            if limitby :
                limitby = limitby.strip().lower()
            if limitby :
                if limitby not in ('quota', 'balance', 'noquota', \
                                            'noprint', 'nochange') :
                    raise PyKotaCommandLineError, _("Invalid limitby value %s") % options["limitby"]
                if limitby in ('noquota', 'nochange') :    
                    options["noquota"] = 1
                if (limitby in ('nochange', 'noprint')) and options["groups"] :    
                    raise PyKotaCommandLineError, _("Invalid limitby value %s") % options["limitby"]
            
            used = options["used"]
            if used :
                used = used.strip()
                try :
                    int(used)
                except ValueError :
                    raise PyKotaCommandLineError, _("Invalid used value %s.") % used
                    
            increase = options["increase"]
            if increase :
                try :
                    increase = int(increase.strip())
                except ValueError :
                    raise PyKotaCommandLineError, _("Invalid increase value %s.") % increase
            
            if not options["noquota"] :
                if options["softlimit"] :
                    try :
                        softlimit = int(options["softlimit"].strip())
                        if softlimit < 0 :
                            raise ValueError
                    except ValueError :    
                        raise PyKotaCommandLineError, _("Invalid softlimit value %s.") % options["softlimit"]
                if options["hardlimit"] :
                    try :
                        hardlimit = int(options["hardlimit"].strip())
                        if hardlimit < 0 :
                            raise ValueError
                    except ValueError :    
                        raise PyKotaCommandLineError, _("Invalid hardlimit value %s.") % options["hardlimit"]
                if (softlimit is not None) and (hardlimit is not None) and (hardlimit < softlimit) :        
                    # error, exchange them
                    self.printInfo(_("Hard limit %i is less than soft limit %i, values will be exchanged.") % (hardlimit, softlimit))
                    (softlimit, hardlimit) = (hardlimit, softlimit)
                
            overcharge = options["overcharge"]
            if overcharge :
                try :
                    overcharge = float(overcharge.strip())
                except (ValueError, AttributeError) :    
                    raise PyKotaCommandLineError, _("Invalid overcharge value %s") % options["overcharge"]
                    
            balance = options["balance"]
            if balance :
                balance = balance.strip()
                try :
                    balancevalue = float(balance)
                except ValueError :    
                    raise PyKotaCommandLineError, _("Invalid balance value %s") % options["balance"]
                
            if options["charge"] :
                try :
                    charges = [float(part) for part in options["charge"].split(',', 1)]
                except ValueError :    
                    raise PyKotaCommandLineError, _("Invalid charge amount value %s") % options["charge"]
                else :    
                    if len(charges) > 2 :
                        charges = charges[:2]
                    if len(charges) != 2 :
                        charges = [charges[0], None]
                        
            if options["ingroups"] :    
                groupnames = [gname.strip() for gname in options["ingroups"].split(',')]
            else :    
                groupnames = []
                
            rejectunknown = self.config.getRejectUnknown()    
            printeradded = 0
            printers = self.storage.getMatchingPrinters(options["printer"])
            if not printers :
                pname = options["printer"]
                if options["add"] and pname :
                    if self.isValidName(pname) :
                        printers = [ self.storage.addPrinter(pname) ]
                        if printers[0].Exists :
                            printeradded = 1
                        else :    
                            raise PyKotaToolError, _("Impossible to add printer %s") % pname
                    else :    
                        raise PyKotaCommandLineError, _("Invalid printer name %s") % pname
                else :
                    raise PyKotaCommandLineError, _("There's no printer matching %s") % pname
            if not names :    
                names = getattr(self.storage, "getAll%ssNames" % suffix)() # all users or groups
                    
            printersgroups = []        
            if options["pgroups"] :        
                printersgroups = self.storage.getMatchingPrinters(options["pgroups"])
                
            if options["prototype"] :    
                protoentry = getattr(self.storage, "get%s" % suffix)(options["prototype"])
                if not protoentry.Exists :
                    raise PyKotaCommandLineError, _("Prototype object %s not found in Quota Storage.") % protoentry.Name
                else :    
                    limitby = protoentry.LimitBy
                    balancevalue = protoentry.AccountBalance
                    if balancevalue is not None :
                        balance = str(abs(balancevalue))
                    else :    
                        balance = None
                    overcharge = getattr(protoentry, "OverCharge", None)
                
            sys.stdout.write(_("Managing print quotas... (this may take a lot of time)"))
            sys.stdout.flush()
            missingusers = {}
            missinggroups = {}    
            changed = {} # tracks changes made at the user/group level
            for printer in printers :
                for pgroup in printersgroups :
                    pgroup.addPrinterToGroup(printer)    
                    
                if options["charge"] :
                    (perpage, perjob) = charges
                    printer.setPrices(perpage, perjob)    
                    
                if options["prototype"] :
                    protoquota = getattr(self.storage, "get%sPQuota" % suffix)(protoentry, printer)
                    if not protoquota.Exists :
                        self.printInfo(_("Prototype %s not found in Quota Storage for printer %s.") % (protoentry.Name, printer.Name))
                    else :    
                        (softlimit, hardlimit) = (protoquota.SoftLimit, protoquota.HardLimit)
                        
                if not options["noquota"] :    
                    if hardlimit is None :    
                        hardlimit = softlimit
                        if hardlimit is not None :
                            self.printInfo(_("Undefined hard limit set to soft limit (%s) on printer %s.") % (str(hardlimit), printer.Name))
                    if softlimit is None :    
                        softlimit = hardlimit
                        if softlimit is not None :
                            self.printInfo(_("Undefined soft limit set to hard limit (%s) on printer %s.") % (str(softlimit), printer.Name))
                            
                if options["add"] :    
                    allentries = []    
                    for name in names :
                        email = ""
                        if not options["groups"] :
                            splitname = name.split('/', 1)     # username/email
                            if len(splitname) == 1 :
                                splitname.append("")
                            (name, email) = splitname
                            if email and (email.count('@') != 1) :
                                self.printInfo(_("Invalid email address %s") % email)
                                email = ""
                        entry = getattr(self.storage, "get%s" % suffix)(name)
                        if email and not options["groups"] :
                            entry.Email = email
                        entrypquota = getattr(self.storage, "get%sPQuota" % suffix)(entry, printer)
                        allentries.append((entry, entrypquota))
                else :   
                    allentries = getattr(self.storage, "getPrinter%ssAndQuotas" % suffix)(printer, names)
                    
                # TODO : do this only once !!!    
                allnames = [entry.Name for (entry, dummy) in allentries]
                for name in names :    
                    if not self.matchString(name, allnames) :
                        if options["groups"] :
                            missinggroups[name] = 1
                        else :    
                            missingusers[name] = 1
                    
                for (entry, entrypquota) in allentries :
                    if not changed.has_key(entry.Name) :
                        changed[entry.Name] = {}
                        if not options["groups"] :
                            changed[entry.Name]["ingroups"] = []
                            
                    if not entry.Exists :        
                        # not found
                        if options["add"] :
                            # In case we want to add something, it is crucial
                            # that we DON'T check with the system accounts files
                            # like /etc/passwd because users may be defined 
                            # only remotely
                            if self.isValidName(entry.Name) :
                                reject = 0
                                if rejectunknown :
                                    if options["groups"] :
                                        try :
                                            grp.getgrnam(entry.Name)
                                        except KeyError :    
                                            self.printInfo(_("Unknown group %s") % entry.Name, "error")
                                            reject = 1
                                    else :    
                                        try :
                                            pwd.getpwnam(entry.Name)
                                        except KeyError :    
                                            self.printInfo(_("Unknown user %s") % entry.Name, "error")
                                            reject = 1
                                if not reject :        
                                    entry = getattr(self.storage, "add%s" % suffix)(entry)
                            else :    
                                if options["groups"] :
                                    self.printInfo(_("Invalid group name %s") % entry.Name)
                                else :    
                                    self.printInfo(_("Invalid user name %s") % entry.Name)
                        else :
                            if options["groups"] :
                                missinggroups[entry.Name] = 1
                            else :    
                                missingusers[entry.Name] = 1
                                    
                    if entry.Exists and (not entrypquota.Exists) :
                        # not found
                        if options["add"] :
                            entrypquota = getattr(self.storage, "add%sPQuota" % suffix)(entry, printer)
                            
                    if not entrypquota.Exists :     
                        self.printInfo(_("Quota not found for object %s on printer %s.") % (entry.Name, printer.Name))
                    else :    
                        if options["noquota"] or options["prototype"] \
                           or ((softlimit is not None) and (hardlimit is not None)) :
                            entrypquota.setLimits(softlimit, hardlimit)
                        if increase :
                           if (entrypquota.SoftLimit is None) \
                               or (entrypquota.HardLimit is None) :
                               self.printInfo(_("You can't increase limits by %s when no limit is set.") % increase, "error")
                           else :
                               newsoft = entrypquota.SoftLimit + increase         
                               newhard = entrypquota.HardLimit + increase         
                               if (newsoft >= 0) and (newhard >= 0) :
                                   entrypquota.setLimits(newsoft, newhard)
                               else :    
                                   self.printInfo(_("You can't set negative limits."), "error")
                        if limitby :
                            if changed[entry.Name].get("limitby") is None :
                                entry.setLimitBy(limitby)
                                changed[entry.Name]["limitby"] = limitby
                        
                        if options["reset"] :
                            entrypquota.reset()
                            
                        if options["hardreset"] :    
                            entrypquota.hardreset()
                            
                        if not options["groups"] :
                            if used :
                                entrypquota.setUsage(used)
                                
                            if overcharge is not None :    
                                if changed[entry.Name].get("overcharge") is None :
                                    entry.setOverChargeFactor(overcharge)
                                    changed[entry.Name]["overcharge"] = overcharge
                                    
                            if balance :
                                if changed[entry.Name].get("balance") is None :
                                    if balance.startswith("+") or balance.startswith("-") :
                                        newbalance = float(entry.AccountBalance or 0.0) + balancevalue
                                        newlifetimepaid = float(entry.LifeTimePaid or 0.0) + balancevalue
                                        entry.setAccountBalance(newbalance, newlifetimepaid, options["comment"])
                                    else :
                                        diff = balancevalue - float(entry.AccountBalance or 0.0)
                                        newlifetimepaid = float(entry.LifeTimePaid or 0.0) + diff
                                        entry.setAccountBalance(balancevalue, newlifetimepaid, options["comment"])
                                    changed[entry.Name]["balance"] = balance
                                    
                            for groupname in groupnames :        
                                # not executed if option --ingroups is not used
                                if groupname not in changed[entry.Name]["ingroups"] :
                                    group = self.storage.getGroup(groupname)
                                    if group.Exists :
                                        self.storage.addUserToGroup(entry, group)
                                        changed[entry.Name]["ingroups"].append(groupname)
                                    else :
                                        self.printInfo(_("Group %s not found in the PyKota Storage.") % groupname)
                            
            # Now outputs the list of nonexistent users and groups                
            for name in missingusers.keys() :
                self.printInfo(_("Nonexistent user %s or missing print quota entry.") % name, level="warn")
            for name in missinggroups.keys() :
                self.printInfo(_("Nonexistent group %s or missing print quota entry.") % name, level="warn")
            
            sys.stdout.write("\nDone.\n")    
                     
if __name__ == "__main__" : 
    retcode = 0
    try :
        defaults = { \
                     "printer" : "*", \
                     "comment" : "", \
                   }
        short_options = "vhdo:c:C:l:b:i:naugrp:P:S:H:G:RU:I:"
        long_options = ["help", "version", "comment=", \
                        "overcharge=", "charge=", "delete", "limitby=", \
                        "balance=", "ingroups=", "noquota", "add", "users", \
                        "groups", "reset", "hardreset", "prototype=", \
                        "printer=", "softlimit=", "hardlimit=", "pgroups=", \
                        "increase=", "used="]
        
        # Initializes the command line tool
        editor = EdPyKota(doc=__doc__)
        editor.deferredInit()
        
        # parse and checks the command line
        (options, args) = editor.parseCommandline(sys.argv[1:], short_options, long_options)
        
        # sets long options
        options["help"] = options["h"] or options["help"]
        options["version"] = options["v"] or options["version"]
        options["add"] = options["a"] or options["add"]
        options["users"] = options["u"] or options["users"]
        options["groups"] = options["g"] or options["groups"]
        options["prototype"] = options["p"] or options["prototype"]
        options["printer"] = options["P"] or options["printer"] or defaults["printer"]
        options["softlimit"] = options["S"] or options["softlimit"]
        options["hardlimit"] = options["H"] or options["hardlimit"] 
        options["reset"] = options["r"] or options["reset"] 
        options["noquota"] = options["n"] or options["noquota"]
        options["limitby"] = options["l"] or options["limitby"]
        options["balance"] = options["b"] or options["balance"] 
        options["delete"] = options["d"] or options["delete"] 
        options["ingroups"] = options["i"] or options["ingroups"]
        options["charge"] = options["c"] or options["charge"]
        options["pgroups"] = options["G"] or options["pgroups"]
        options["hardreset"] = options["R"] or options["hardreset"] 
        options["used"] = options["U"] or options["used"]
        options["overcharge"] = options["o"] or options["overcharge"]
        options["comment"] = options["C"] or options["comment"] or defaults["comment"]
        options["increase"] = options["I"] or options["increase"]
        
        if options["help"] :
            editor.display_usage_and_quit()
        elif options["version"] :
            editor.display_version_and_quit()
        elif options["users"] and options["groups"] :    
            raise PyKotaCommandLineError, _("incompatible options, see help.")
        elif (options["add"] or options["prototype"]) and options["delete"] :    
            raise PyKotaCommandLineError, _("incompatible options, see help.")
        elif (options["reset"] or options["hardreset"] or options["limitby"] or options["used"] or options["balance"] or options["overcharge"] or options["softlimit"] or options["hardlimit"]) and options["prototype"] :
            raise PyKotaCommandLineError, _("incompatible options, see help.")
        elif options["noquota"] and (options["prototype"] or options["hardlimit"] or options["softlimit"]) :
            raise PyKotaCommandLineError, _("incompatible options, see help.")
        elif options["groups"] and (options["balance"] or options["ingroups"] or options["used"] or options["overcharge"]) :
            raise PyKotaCommandLineError, _("incompatible options, see help.")
        elif options["comment"] and not options["balance"] :    
            raise PyKotaCommandLineError, _("incompatible options, see help.")
        else :
            retcode = editor.main(args, options)
    except KeyboardInterrupt :        
        sys.stderr.write("\nInterrupted with Ctrl+C !\n")
        retcode = -3
    except PyKotaCommandLineError, msg :     
        sys.stderr.write("%s : %s\n" % (sys.argv[0], msg))
        retcode = -2
    except SystemExit :        
        pass
    except :
        try :
            editor.crashed("edpykota failed")
        except :    
            crashed("edpykota failed")
        retcode = -1

    try :
        editor.storage.close()
    except (TypeError, NameError, AttributeError) :    
        pass
        
    sys.exit(retcode)