imip-agent

tools/update_quotas.py

1061:5f9f96eb4f8f
2016-02-09 Paul Boddie Updated the administration page, removing obsolete settings references.
     1 #!/usr/bin/env python     2      3 """     4 Remove expired events from quota journals.     5      6 Copyright (C) 2014, 2015, 2016 Paul Boddie <paul@boddie.org.uk>     7      8 This program is free software; you can redistribute it and/or modify it under     9 the terms of the GNU General Public License as published by the Free Software    10 Foundation; either version 3 of the License, or (at your option) any later    11 version.    12     13 This program is distributed in the hope that it will be useful, but WITHOUT    14 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS    15 FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more    16 details.    17     18 You should have received a copy of the GNU General Public License along with    19 this program.  If not, see <http://www.gnu.org/licenses/>.    20 """    21     22 from os.path import split    23 import sys    24     25 # Find the modules.    26     27 try:    28     import imiptools    29 except ImportError:    30     parent = split(split(__file__)[0])[0]    31     if split(parent)[1] == "imip-agent":    32         sys.path.append(parent)    33     34 from codecs import getwriter    35 from imiptools.dates import get_datetime, get_default_timezone, get_time, \    36                             to_utc_datetime    37 from imip_store import FileJournal    38     39 def remove_expired_entries(entries, expiry):    40     41     "Remove from 'entries' events that end at or before 'expiry'."    42     43     removed = []    44     i = 0    45     46     while i < len(entries):    47         period = entries[i]    48     49         if period.get_end_point() <= expiry:    50             removed.append(period)    51             del entries[i]    52         else:    53             i += 1    54     55     return removed    56     57 def update_entries(journal, quota, expiry, store, verbose):    58     59     """    60     Using the given 'journal' process quota records for the given 'quota', with    61     the given 'expiry' time used to expire events ending before or at this time,    62     with None meaning the current time.    63     64     If 'store' is set, the stored details will be updated; otherwise, the    65     details will be written to standard output.    66     67     If 'verbose' is set, messages will be written to standard error.    68     """    69     70     if not store:    71         stdout = getwriter("utf-8")(sys.stdout)    72     if verbose:    73         stderr = getwriter("utf-8")(sys.stderr)    74     75     if not expiry:    76         expiry = get_time()    77     78     journal.acquire_lock(quota)    79     80     try:    81         for user in journal.get_quota_users(quota):    82             if verbose:    83                 print >>stderr, user    84     85             entries = journal.get_entries(quota, user)    86             removed = remove_expired_entries(entries, expiry)    87     88             if verbose:    89                 for period in removed:    90                     print >>stderr, "\t".join(("Removed",) + period.as_tuple(strings_only=True))    91     92             # Store the processed entries.    93     94             if store:    95                 journal.set_entries(quota, user, entries)    96     97             # Alternatively, just write the entries to standard output.    98     99             else:   100                 for period in entries:   101                     print >>stdout, "\t".join(period.as_tuple(strings_only=True))   102     finally:   103         journal.release_lock(quota)   104    105 # Main program.   106    107 if __name__ == "__main__":   108    109     # Interpret the command line arguments.   110    111     quotas = []   112     args = []   113     journal_dir = []   114     expiry = []   115     ignored = []   116    117     # Collect quota details first, switching to other arguments when encountering   118     # switches.   119    120     l = quotas   121    122     for arg in sys.argv[1:]:   123         if arg in ("-s", "-v"):   124             args.append(arg)   125             l = ignored   126         elif arg == "-j":   127             l = journal_dir   128         elif arg == "-e":   129             l = expiry   130         else:   131             l.append(arg)   132    133     try:   134         quota = quotas[0]   135     except IndexError:   136         print >>sys.stderr, """\   137 Usage: %s <quota> <options>   138    139 Need a quota along with the -s option if updating the journal.   140 Specify -v for additional messages on standard error.   141    142 General options:   143    144 -e  indicate an expiry time for events (default is now)   145 -j  indicate the journal directory location   146 """ % split(sys.argv[0])[1]   147         sys.exit(1)   148    149     # Define any other options.   150    151     store = "-s" in args   152     verbose = "-v" in args   153    154     # Override defaults if indicated.   155    156     journal_dir = journal_dir and journal_dir[0] or None   157     expiry = expiry and expiry[0] or None   158    159     if expiry:   160         expiry = to_utc_datetime(get_datetime(expiry), get_default_timezone())   161         if not expiry:   162             print >>sys.stderr, "Expiry time must be a valid datetime."   163             sys.exit(1)   164    165     # Obtain store-related objects.   166    167     journal = FileJournal(journal_dir)   168    169     # Obtain a list of users for processing.   170    171     if quota in ("*", "all"):   172         quotas = journal.get_quotas()   173    174     # Process the given users.   175    176     if verbose:   177         stderr = getwriter("utf-8")(sys.stderr)   178    179     for quota in quotas:   180         if verbose:   181             print >>stderr, quota   182         update_entries(journal, quota, expiry, store, verbose)   183    184 # vim: tabstop=4 expandtab shiftwidth=4