1 #!/usr/bin/env python 2 3 """ 4 Remove expired events from quota journals. 5 6 Copyright (C) 2014, 2015, 2016, 2017 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 abspath, split 23 import sys 24 25 # Find the modules. 26 27 try: 28 import imiptools 29 except ImportError: 30 parent = abspath(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.config import settings 36 from imiptools.dates import get_datetime, get_default_timezone, get_time, \ 37 to_utc_datetime 38 from imiptools.stores import get_journal 39 40 def remove_expired_entries(entries, expiry): 41 42 "Remove from 'entries' events that end at or before 'expiry'." 43 44 removed = [] 45 i = 0 46 47 while i < len(entries): 48 period = entries[i] 49 50 if period.get_end_point() <= expiry: 51 removed.append(period) 52 del entries[i] 53 else: 54 i += 1 55 56 return removed 57 58 def update_entries(journal, quota, expiry, store, verbose): 59 60 """ 61 Using the given 'journal' process quota records for the given 'quota', with 62 the given 'expiry' time used to expire events ending before or at this time, 63 with None meaning the current time. 64 65 If 'store' is set, the stored details will be updated; otherwise, the 66 details will be written to standard output. 67 68 If 'verbose' is set, messages will be written to standard error. 69 """ 70 71 if not store: 72 stdout = getwriter("utf-8")(sys.stdout) 73 if verbose: 74 stderr = getwriter("utf-8")(sys.stderr) 75 76 if not expiry: 77 expiry = get_time() 78 79 journal.acquire_lock(quota) 80 81 try: 82 for user in journal.get_quota_users(quota): 83 if verbose: 84 print >>stderr, user 85 86 entries = journal.get_entries(quota, user) 87 removed = remove_expired_entries(entries, expiry) 88 89 if verbose: 90 for period in removed: 91 print >>stderr, "\t".join(("Removed",) + period.as_tuple(strings_only=True)) 92 93 # Store the processed entries. 94 95 if store: 96 journal.set_entries(quota, user, entries) 97 98 # Alternatively, just write the entries to standard output. 99 100 else: 101 for period in entries: 102 print >>stdout, "\t".join(period.as_tuple(strings_only=True)) 103 finally: 104 journal.release_lock(quota) 105 106 # Main program. 107 108 if __name__ == "__main__": 109 110 # Interpret the command line arguments. 111 112 quotas = [] 113 args = [] 114 store_type = [] 115 journal_dir = [] 116 expiry = [] 117 ignored = [] 118 119 # Collect quota details first, switching to other arguments when encountering 120 # switches. 121 122 l = quotas 123 124 for arg in sys.argv[1:]: 125 if arg in ("-s", "-v"): 126 args.append(arg) 127 l = ignored 128 elif arg == "-T": 129 l = store_type 130 elif arg == "-j": 131 l = journal_dir 132 elif arg == "-e": 133 l = expiry 134 else: 135 l.append(arg) 136 137 try: 138 quota = quotas[0] 139 except IndexError: 140 print >>sys.stderr, """\ 141 Usage: %s <quota> <options> 142 143 Need a quota along with the -s option if updating the journal. 144 Specify -v for additional messages on standard error. 145 146 General options: 147 148 -e Indicates an expiry time for events (default is now) 149 -j Indicates the journal directory location 150 -T Indicates the store type (the configured value if omitted) 151 """ % split(sys.argv[0])[1] 152 sys.exit(1) 153 154 # Define any other options. 155 156 store = "-s" in args 157 verbose = "-v" in args 158 159 # Override defaults if indicated. 160 161 getvalue = lambda value, default=None: value and value[0] or default 162 163 store_type = getvalue(store_type, settings["STORE_TYPE"]) 164 journal_dir = getvalue(journal_dir) 165 expiry = getvalue(expiry) 166 167 if expiry: 168 expiry = to_utc_datetime(get_datetime(expiry), get_default_timezone()) 169 if not expiry: 170 print >>sys.stderr, "Expiry time must be a valid datetime." 171 sys.exit(1) 172 173 # Obtain store-related objects. 174 175 journal = get_journal(store_type, journal_dir) 176 177 # Obtain a list of users for processing. 178 179 if quota in ("*", "all"): 180 quotas = journal.get_quotas() 181 182 # Process the given users. 183 184 if verbose: 185 stderr = getwriter("utf-8")(sys.stderr) 186 187 for quota in quotas: 188 if verbose: 189 print >>stderr, quota 190 update_entries(journal, quota, expiry, store, verbose) 191 192 # vim: tabstop=4 expandtab shiftwidth=4