summaryrefslogtreecommitdiffstats
path: root/testing/mozharness/scripts/release/generate-checksums.py
blob: 6d01923e444c816d7f553f225d991c2a8720fbb3 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

import binascii
import hashlib
import os
import re
import sys
from multiprocessing.pool import ThreadPool

import six

sys.path.insert(1, os.path.dirname(os.path.dirname(sys.path[0])))

from mozharness.base.python import VirtualenvMixin, virtualenv_config_options
from mozharness.base.script import BaseScript
from mozharness.mozilla.checksums import parse_checksums_file
from mozharness.mozilla.merkle import MerkleTree


class ChecksumsGenerator(BaseScript, VirtualenvMixin):
    config_options = [
        [
            ["--stage-product"],
            {
                "dest": "stage_product",
                "help": "Name of product used in file server's directory structure, "
                "e.g.: firefox, mobile",
            },
        ],
        [
            ["--version"],
            {
                "dest": "version",
                "help": "Version of release, e.g.: 59.0b5",
            },
        ],
        [
            ["--build-number"],
            {
                "dest": "build_number",
                "help": "Build number of release, e.g.: 2",
            },
        ],
        [
            ["--bucket-name"],
            {
                "dest": "bucket_name",
                "help": "Full bucket name e.g.: net-mozaws-prod-delivery-{firefox,archive}.",
            },
        ],
        [
            ["-j", "--parallelization"],
            {
                "dest": "parallelization",
                "default": 20,
                "type": int,
                "help": "Number of checksums file to download concurrently",
            },
        ],
        [
            ["--branch"],
            {
                "dest": "branch",
                "help": "dummy option",
            },
        ],
        [
            ["--build-pool"],
            {
                "dest": "build_pool",
                "help": "dummy option",
            },
        ],
    ] + virtualenv_config_options

    def __init__(self):
        BaseScript.__init__(
            self,
            config_options=self.config_options,
            require_config_file=False,
            config={
                "virtualenv_modules": [
                    "boto",
                ],
                "virtualenv_path": "venv",
            },
            all_actions=[
                "create-virtualenv",
                "collect-individual-checksums",
                "create-big-checksums",
                "create-summary",
            ],
            default_actions=[
                "create-virtualenv",
                "collect-individual-checksums",
                "create-big-checksums",
                "create-summary",
            ],
        )

        self.checksums = {}
        self.file_prefix = self._get_file_prefix()

    def _pre_config_lock(self, rw_config):
        super(ChecksumsGenerator, self)._pre_config_lock(rw_config)

        # These defaults are set here rather in the config because default
        # lists cannot be completely overidden, only appended to.
        if not self.config.get("formats"):
            self.config["formats"] = ["sha512", "sha256"]

        if not self.config.get("includes"):
            self.config["includes"] = [
                r"^.*\.tar\.bz2$",
                r"^.*\.tar\.xz$",
                r"^.*\.snap$",
                r"^.*\.dmg$",
                r"^.*\.pkg$",
                r"^.*\.bundle$",
                r"^.*\.mar$",
                r"^.*Setup.*\.exe$",
                r"^.*Installer\.exe$",
                r"^.*\.msi$",
                r"^.*\.xpi$",
                r"^.*fennec.*\.apk$",
                r"^.*/jsshell.*$",
            ]

    def _get_file_prefix(self):
        return "pub/{}/candidates/{}-candidates/build{}/".format(
            self.config["stage_product"],
            self.config["version"],
            self.config["build_number"],
        )

    def _get_sums_filename(self, format_):
        return "{}SUMS".format(format_.upper())

    def _get_summary_filename(self, format_):
        return "{}SUMMARY".format(format_.upper())

    def _get_hash_function(self, format_):
        if format_ in ("sha256", "sha384", "sha512"):
            return getattr(hashlib, format_)
        else:
            self.fatal("Unsupported format {}".format(format_))

    def _get_bucket(self):
        self.activate_virtualenv()
        from boto import connect_s3

        self.info("Connecting to S3")
        conn = connect_s3(anon=True)
        self.info("Connecting to bucket {}".format(self.config["bucket_name"]))
        self.bucket = conn.get_bucket(self.config["bucket_name"])
        return self.bucket

    def collect_individual_checksums(self):
        """This step grabs all of the small checksums files for the release,
        filters out any unwanted files from within them, and adds the remainder
        to self.checksums for subsequent steps to use."""
        bucket = self._get_bucket()
        self.info("File prefix is: {}".format(self.file_prefix))

        # temporary holding place for checksums
        raw_checksums = []

        def worker(item):
            self.debug("Downloading {}".format(item))
            sums = bucket.get_key(item).get_contents_as_string()
            raw_checksums.append(sums)

        def find_checksums_files():
            self.info("Getting key names from bucket")
            checksum_files = {"beets": [], "checksums": []}
            for key in bucket.list(prefix=self.file_prefix):
                if key.key.endswith(".checksums"):
                    self.debug("Found checksums file: {}".format(key.key))
                    checksum_files["checksums"].append(key.key)
                elif key.key.endswith(".beet"):
                    self.debug("Found beet file: {}".format(key.key))
                    checksum_files["beets"].append(key.key)
                else:
                    self.debug("Ignoring non-checksums file: {}".format(key.key))
            if checksum_files["beets"]:
                self.log("Using beet format")
                return checksum_files["beets"]
            else:
                self.log("Using checksums format")
                return checksum_files["checksums"]

        pool = ThreadPool(self.config["parallelization"])
        pool.map(worker, find_checksums_files())

        for c in raw_checksums:
            for f, info in six.iteritems(parse_checksums_file(c)):
                for pattern in self.config["includes"]:
                    if re.search(pattern, f):
                        if f in self.checksums:
                            if info == self.checksums[f]:
                                self.debug(
                                    "Duplicate checksum for file {}"
                                    " but the data matches;"
                                    " continuing...".format(f)
                                )
                                continue
                            self.fatal(
                                "Found duplicate checksum entry for {}, "
                                "don't know which one to pick.".format(f)
                            )
                        if not set(self.config["formats"]) <= set(info["hashes"]):
                            self.fatal("Missing necessary format for file {}".format(f))
                        self.debug("Adding checksums for file: {}".format(f))
                        self.checksums[f] = info
                        break
                else:
                    self.debug("Ignoring checksums for file: {}".format(f))

    def create_summary(self):
        """
        This step computes a Merkle tree over the checksums for each format
        and writes a file containing the head of the tree and inclusion proofs
        for each file.
        """
        for fmt in self.config["formats"]:
            hash_fn = self._get_hash_function(fmt)
            files = [fn for fn in sorted(self.checksums)]
            data = [self.checksums[fn]["hashes"][fmt] for fn in files]

            tree = MerkleTree(hash_fn, data)
            head = binascii.hexlify(tree.head())
            proofs = [
                binascii.hexlify(tree.inclusion_proof(i).to_rfc6962_bis())
                for i in range(len(files))
            ]

            summary = self._get_summary_filename(fmt)
            self.info("Creating summary file: {}".format(summary))

            content = "{} TREE_HEAD\n".format(head.decode("ascii"))
            for i in range(len(files)):
                content += "{} {}\n".format(proofs[i].decode("ascii"), files[i])

            self.write_to_file(summary, content)

    def create_big_checksums(self):
        for fmt in self.config["formats"]:
            sums = self._get_sums_filename(fmt)
            self.info("Creating big checksums file: {}".format(sums))
            with open(sums, "w+") as output_file:
                for fn in sorted(self.checksums):
                    output_file.write(
                        "{}  {}\n".format(
                            self.checksums[fn]["hashes"][fmt].decode("ascii"), fn
                        )
                    )


if __name__ == "__main__":
    myScript = ChecksumsGenerator()
    myScript.run_and_exit()