← Back to branch summary

~launchpad-pqm/launchpad/devel

« back to all changes in this revision

Viewing changes to cronscripts/update-cve.py

Merged mainline.

Show diffs side-by-side

added added

removed removed

Lines of Context:
cronscripts/update-cve.py
16
16
 
17
17
import cElementTree
18
18
 
 
19
from zope.interface import implements
 
20
 
19
21
from canonical.config import config
 
22
from canonical.launchpad.interfaces.looptuner import ITunableLoop
20
23
from canonical.launchpad.scripts.cveimport import CVEDB_NS, update_one_cve
21
 
 
22
24
from canonical.launchpad.scripts.base import (
23
25
    LaunchpadCronScript, LaunchpadScriptFailure)
 
26
from canonical.launchpad.utilities.looptuner import LoopTuner
 
27
 
 
28
 
 
29
class CveUpdaterTunableLoop(object):
 
30
    """An `ITunableLoop` for updating CVEs."""
 
31
 
 
32
    implements(ITunableLoop)
 
33
 
 
34
    total_updated = 0
 
35
 
 
36
    def __init__(self, cves, transaction, logger, offset=0):
 
37
        self.cves = cves
 
38
        self.transaction = transaction
 
39
        self.logger = logger
 
40
        self.offset = offset
 
41
        self.total_updated = 0
 
42
 
 
43
    def isDone(self):
 
44
        """See `ITunableLoop`."""
 
45
        return self.offset is None
 
46
 
 
47
    def __call__(self, chunk_size):
 
48
        """Retrieve a batch of CVEs and update them.
 
49
 
 
50
        See `ITunableLoop`.
 
51
        """
 
52
        chunk_size = int(chunk_size)
 
53
 
 
54
        self.logger.debug("More %d" % chunk_size)
 
55
 
 
56
        start = self.offset
 
57
        end = self.offset + chunk_size
 
58
 
 
59
        self.transaction.begin()
 
60
 
 
61
        cve_batch = self.cves[start:end]
 
62
        self.offset = None
 
63
        for cve in cve_batch:
 
64
            start += 1
 
65
            self.offset = start
 
66
            update_one_cve(cve, self.logger)
 
67
            self.total_updated += 1
 
68
 
 
69
        self.logger.debug("Committing.")
 
70
        self.transaction.commit()
24
71
 
25
72
 
26
73
class CVEUpdater(LaunchpadCronScript):
 
74
 
27
75
    def add_my_options(self):
28
76
        """Parse command line arguments."""
29
 
        self.parser.add_option("-f", "--cvefile", dest="cvefile",
30
 
                               default=None,
31
 
                               help="An XML file containing the CVE database.")
32
 
        self.parser.add_option("-u", "--cveurl", dest="cveurl",
33
 
                               default=config.cveupdater.cve_db_url,
34
 
                               help="The URL for the gzipped XML CVE database.")
 
77
        self.parser.add_option(
 
78
            "-f", "--cvefile", dest="cvefile", default=None,
 
79
            help="An XML file containing the CVE database.")
 
80
        self.parser.add_option(
 
81
            "-u", "--cveurl", dest="cveurl",
 
82
            default=config.cveupdater.cve_db_url,
 
83
            help="The URL for the gzipped XML CVE database.")
35
84
 
36
85
    def main(self):
37
86
        self.logger.info('Initializing...')
55
104
 
56
105
            cve_db_gz = url.read()
57
106
            self.logger.info("%d bytes downloaded." % len(cve_db_gz))
58
 
            cve_db = gzip.GzipFile(fileobj=StringIO.StringIO(cve_db_gz)).read()
 
107
            cve_db = gzip.GzipFile(
 
108
                fileobj=StringIO.StringIO(cve_db_gz)).read()
59
109
        else:
60
110
            raise LaunchpadScriptFailure('No CVE database file or URL given.')
61
111
 
65
115
        dom = cElementTree.fromstring(cve_db)
66
116
        items = dom.findall(CVEDB_NS + 'item')
67
117
        self.logger.info("Updating database...")
68
 
        for item in items:
69
 
            self.txn.begin()
70
 
            update_one_cve(item, self.logger)
71
 
            self.txn.commit()
 
118
 
 
119
        # We use Looptuner to control the ideal number of CVEs
 
120
        # processed in each transaction, during at least 2 seconds.
 
121
        loop = CveUpdaterTunableLoop(items, self.txn, self.logger)
 
122
        loop_tuner = LoopTuner(loop, 2)
 
123
        loop_tuner.run()
 
124
 
72
125
        timing.finish()
73
126
        self.logger.info('%d seconds to update database.' % timing.seconds())
74
127