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 imiptools.stores.file 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