1#!/usr/bin/env python3
2
3import contextlib
4import hashlib
5import os
6import platform
7import re
8import shutil
9import subprocess
10import sys
11import tarfile
12import urllib.request
13import zipfile
14from dataclasses import dataclass
15from pathlib import Path
16from typing import Optional, Union
17
18import click
19import cryptography.fernet
20import parver
21
22
23@contextlib.contextmanager
24def chdir(path: Path):  # pragma: no cover
25    old_dir = os.getcwd()
26    os.chdir(path)
27    yield
28    os.chdir(old_dir)
29
30
31class BuildError(Exception):
32    pass
33
34
35def bool_from_env(envvar: str) -> bool:
36    val = os.environ.get(envvar, "")
37    if not val or val.lower() in ("0", "false"):
38        return False
39    else:
40        return True
41
42
43class ZipFile2(zipfile.ZipFile):
44    # ZipFile and tarfile have slightly different APIs. Let's fix that.
45    def add(self, name: str, arcname: str) -> None:
46        return self.write(name, arcname)
47
48    def __enter__(self) -> "ZipFile2":
49        return self
50
51
52@dataclass(frozen=True, repr=False)
53class BuildEnviron:
54    PLATFORM_TAGS = {
55        "Darwin": "osx",
56        "Windows": "windows",
57        "Linux": "linux",
58    }
59
60    system: str
61    root_dir: Path
62    branch: Optional[str] = None
63    tag: Optional[str] = None
64    is_pull_request: bool = True
65    should_build_wheel: bool = False
66    should_build_docker: bool = False
67    should_build_pyinstaller: bool = False
68    should_build_wininstaller: bool = False
69    has_aws_creds: bool = False
70    has_twine_creds: bool = False
71    docker_username: Optional[str] = None
72    docker_password: Optional[str] = None
73    build_key: Optional[str] = None
74
75    @classmethod
76    def from_env(cls) -> "BuildEnviron":
77        branch = None
78        tag = None
79
80        if ref := os.environ.get("GITHUB_REF", ""):
81            if ref.startswith("refs/heads/"):
82                branch = ref.replace("refs/heads/", "")
83            if ref.startswith("refs/pull/"):
84                branch = "pr-" + ref.split("/")[2]
85            if ref.startswith("refs/tags/"):
86                tag = ref.replace("refs/tags/", "")
87
88        is_pull_request = os.environ.get("GITHUB_EVENT_NAME", "pull_request") == "pull_request"
89
90        return cls(
91            system=platform.system(),
92            root_dir=Path(__file__).parent.parent,
93            branch=branch,
94            tag=tag,
95            is_pull_request=is_pull_request,
96            should_build_wheel=bool_from_env("CI_BUILD_WHEEL"),
97            should_build_pyinstaller=bool_from_env("CI_BUILD_PYINSTALLER"),
98            should_build_wininstaller=bool_from_env("CI_BUILD_WININSTALLER"),
99            should_build_docker=bool_from_env("CI_BUILD_DOCKER"),
100            has_aws_creds=bool_from_env("AWS_ACCESS_KEY_ID"),
101            has_twine_creds=bool_from_env("TWINE_USERNAME") and bool_from_env("TWINE_PASSWORD"),
102            docker_username=os.environ.get("DOCKER_USERNAME", None),
103            docker_password=os.environ.get("DOCKER_PASSWORD", None),
104            build_key=os.environ.get("CI_BUILD_KEY", None),
105        )
106
107    def archive(self, path: Path) -> Union[tarfile.TarFile, ZipFile2]:
108        if self.system == "Windows":
109            return ZipFile2(path, "w")
110        else:
111            return tarfile.open(path, "w:gz")
112
113    @property
114    def archive_path(self) -> Path:
115        if self.system == "Windows":
116            ext = "zip"
117        else:
118            ext = "tar.gz"
119        return self.dist_dir / f"mitmproxy-{self.version}-{self.platform_tag}.{ext}"
120
121    @property
122    def build_dir(self) -> Path:
123        return self.release_dir / "build"
124
125    @property
126    def dist_dir(self) -> Path:
127        return self.release_dir / "dist"
128
129    @property
130    def docker_tag(self) -> str:
131        if self.branch == "main":
132            t = "dev"
133        else:
134            t = self.version
135        return f"mitmproxy/mitmproxy:{t}"
136
137    def dump_info(self, fp=sys.stdout) -> None:
138        lst = [
139            "version",
140            "tag",
141            "branch",
142            "platform_tag",
143            "root_dir",
144            "release_dir",
145            "build_dir",
146            "dist_dir",
147            "upload_dir",
148            "should_build_wheel",
149            "should_build_pyinstaller",
150            "should_build_wininstaller",
151            "should_build_docker",
152            "should_upload_aws",
153            "should_upload_docker",
154            "should_upload_pypi",
155        ]
156        for attr in lst:
157            print(f"cibuild.{attr}={getattr(self, attr)}", file=fp)
158
159    def check_version(self) -> None:
160        """
161        Check that version numbers match our conventions.
162        Raises a ValueError if there is a mismatch.
163        """
164        contents = (self.root_dir / "mitmproxy" / "version.py").read_text("utf8")
165        match = re.search(r'^VERSION = "(.+?)"', contents, re.M)
166        assert match
167        version = match.group(1)
168
169        if self.is_prod_release:
170            # For production releases, we require strict version equality
171            if self.version != version:
172                raise ValueError(f"Tag is {self.tag}, but mitmproxy/version.py is {version}.")
173        elif not self.is_maintenance_branch:
174            # Commits on maintenance branches don't need the dev suffix. This
175            # allows us to incorporate and test commits between tagged releases.
176            # For snapshots, we only ensure that mitmproxy/version.py contains a
177            # dev release.
178            version_info = parver.Version.parse(version)
179            if not version_info.is_devrelease:
180                raise ValueError(f"Non-production releases must have dev suffix: {version}")
181
182    @property
183    def is_maintenance_branch(self) -> bool:
184        """
185            Is this an untagged commit on a maintenance branch?
186        """
187        if not self.tag and self.branch and re.match(r"v\d+\.x", self.branch):
188            return True
189        return False
190
191    @property
192    def has_docker_creds(self) -> bool:
193        return bool(self.docker_username and self.docker_password)
194
195    @property
196    def is_prod_release(self) -> bool:
197        if not self.tag or not self.tag.startswith("v"):
198            return False
199        try:
200            v = parver.Version.parse(self.version, strict=True)
201        except (parver.ParseError, BuildError):
202            return False
203        return not v.is_prerelease
204
205    @property
206    def platform_tag(self) -> str:
207        if self.system in self.PLATFORM_TAGS:
208            return self.PLATFORM_TAGS[self.system]
209        raise BuildError(f"Unsupported platform: {self.system}")
210
211    @property
212    def release_dir(self) -> Path:
213        return self.root_dir / "release"
214
215    @property
216    def should_upload_docker(self) -> bool:
217        return all([
218            (self.is_prod_release or self.branch in ["main", "dockertest"]),
219            self.should_build_docker,
220            self.has_docker_creds,
221        ])
222
223    @property
224    def should_upload_aws(self) -> bool:
225        return all([
226            self.has_aws_creds,
227            (self.should_build_wheel or self.should_build_pyinstaller or self.should_build_wininstaller),
228        ])
229
230    @property
231    def should_upload_pypi(self) -> bool:
232        return all([
233            self.is_prod_release,
234            self.should_build_wheel,
235            self.has_twine_creds,
236        ])
237
238    @property
239    def upload_dir(self) -> str:
240        if self.tag:
241            return self.version
242        else:
243            return f"branches/{self.version}"
244
245    @property
246    def version(self) -> str:
247        if self.tag:
248            if self.tag.startswith("v"):
249                try:
250                    parver.Version.parse(self.tag[1:], strict=True)
251                except parver.ParseError:
252                    return self.tag
253                return self.tag[1:]
254            return self.tag
255        elif self.branch:
256            return self.branch
257        else:
258            raise BuildError("We're on neither a tag nor a branch - could not establish version")
259
260
261def build_wheel(be: BuildEnviron) -> None:  # pragma: no cover
262    click.echo("Building wheel...")
263    subprocess.check_call([
264        "python",
265        "setup.py",
266        "-q",
267        "bdist_wheel",
268        "--dist-dir", be.dist_dir,
269    ])
270    whl, = be.dist_dir.glob('mitmproxy-*-py3-none-any.whl')
271    click.echo(f"Found wheel package: {whl}")
272    subprocess.check_call(["tox", "-e", "wheeltest", "--", whl])
273
274
275DOCKER_PLATFORMS = "linux/amd64,linux/arm64"
276
277
278def build_docker_image(be: BuildEnviron) -> None:  # pragma: no cover
279    click.echo("Building Docker images...")
280
281    whl, = be.dist_dir.glob('mitmproxy-*-py3-none-any.whl')
282    docker_build_dir = be.release_dir / "docker"
283    shutil.copy(whl, docker_build_dir / whl.name)
284
285    subprocess.check_call([
286        "docker", "buildx", "build",
287        "--tag", be.docker_tag,
288        "--platform", DOCKER_PLATFORMS,
289        "--build-arg", f"MITMPROXY_WHEEL={whl.name}",
290        "."
291    ], cwd=docker_build_dir)
292    # smoke-test the newly built docker image
293
294    # build again without --platform but with --load to make the tag available,
295    # see https://github.com/docker/buildx/issues/59#issuecomment-616050491
296    subprocess.check_call([
297        "docker", "buildx", "build",
298        "--tag", be.docker_tag,
299        "--load",
300        "--build-arg", f"MITMPROXY_WHEEL={whl.name}",
301        "."
302    ], cwd=docker_build_dir)
303    r = subprocess.run([
304        "docker",
305        "run",
306        "--rm",
307        be.docker_tag,
308        "mitmdump",
309        "--version",
310    ], check=True, capture_output=True)
311    print(r.stdout.decode())
312    assert "Mitmproxy: " in r.stdout.decode()
313
314
315def build_pyinstaller(be: BuildEnviron) -> None:  # pragma: no cover
316    click.echo("Building pyinstaller package...")
317
318    PYINSTALLER_SPEC = be.release_dir / "specs"
319    PYINSTALLER_HOOKS = be.release_dir / "hooks"
320    PYINSTALLER_TEMP = be.build_dir / "pyinstaller"
321    PYINSTALLER_DIST = be.build_dir / "binaries" / be.platform_tag
322
323    if PYINSTALLER_TEMP.exists():
324        shutil.rmtree(PYINSTALLER_TEMP)
325    if PYINSTALLER_DIST.exists():
326        shutil.rmtree(PYINSTALLER_DIST)
327
328    if be.platform_tag == "windows":
329        with chdir(PYINSTALLER_SPEC):
330            click.echo("Building PyInstaller binaries in directory mode...")
331            subprocess.check_call(
332                [
333                    "pyinstaller",
334                    "--clean",
335                    "--workpath", PYINSTALLER_TEMP,
336                    "--distpath", PYINSTALLER_DIST,
337                    "./windows-dir.spec"
338                ]
339            )
340            for tool in ["mitmproxy", "mitmdump", "mitmweb"]:
341                click.echo(f"> {tool} --version")
342                executable = (PYINSTALLER_DIST / "onedir" / tool).with_suffix(".exe")
343                click.echo(subprocess.check_output([executable, "--version"]).decode())
344
345    with be.archive(be.archive_path) as archive:
346        for tool in ["mitmproxy", "mitmdump", "mitmweb"]:
347            # We can't have a folder and a file with the same name.
348            if tool == "mitmproxy":
349                tool = "mitmproxy_main"
350            # Make sure that we are in the spec folder.
351            with chdir(PYINSTALLER_SPEC):
352                click.echo(f"Building PyInstaller {tool} binary...")
353                excludes = []
354                if tool != "mitmweb":
355                    excludes.append("mitmproxy.tools.web")
356                if tool != "mitmproxy_main":
357                    excludes.append("mitmproxy.tools.console")
358
359                subprocess.check_call(
360                    [   # type: ignore
361                        "pyinstaller",
362                        "--clean",
363                        "--workpath", PYINSTALLER_TEMP,
364                        "--distpath", PYINSTALLER_DIST,
365                        "--additional-hooks-dir", PYINSTALLER_HOOKS,
366                        "--onefile",
367                        "--console",
368                        "--icon", "icon.ico",
369                    ]
370                    + [x for e in excludes for x in ["--exclude-module", e]]
371                    + [tool]
372                )
373                # Delete the spec file - we're good without.
374                os.remove(f"{tool}.spec")
375
376            executable = PYINSTALLER_DIST / tool
377            if be.platform_tag == "windows":
378                executable = executable.with_suffix(".exe")
379
380            # Remove _main suffix from mitmproxy executable
381            if "_main" in executable.name:
382                executable = executable.rename(
383                    executable.with_name(executable.name.replace("_main", ""))
384                )
385
386            # Test if it works at all O:-)
387            click.echo(f"> {executable} --version")
388            click.echo(subprocess.check_output([executable, "--version"]).decode())
389
390            archive.add(str(executable), str(executable.name))
391    click.echo("Packed {}.".format(be.archive_path.name))
392
393
394def build_wininstaller(be: BuildEnviron) -> None:  # pragma: no cover
395    click.echo("Building wininstaller package...")
396
397    IB_VERSION = "21.6.0"
398    IB_SETUP_SHA256 = "2bc9f9945cb727ad176aa31fa2fa5a8c57a975bad879c169b93e312af9d05814"
399    IB_DIR = be.release_dir / "installbuilder"
400    IB_SETUP = IB_DIR / "setup" / f"{IB_VERSION}-installer.exe"
401    IB_CLI = Path(fr"C:\Program Files\VMware InstallBuilder Enterprise {IB_VERSION}\bin\builder-cli.exe")
402    IB_LICENSE = IB_DIR / "license.xml"
403
404    if not IB_LICENSE.exists() and not be.build_key:
405        click.echo("Cannot build windows installer without secret key.")
406        return
407
408    if not IB_CLI.exists():
409        if not IB_SETUP.exists():
410            click.echo("Downloading InstallBuilder...")
411
412            def report(block, blocksize, total):
413                done = block * blocksize
414                if round(100 * done / total) != round(100 * (done - blocksize) / total):
415                    click.secho(f"Downloading... {round(100 * done / total)}%")
416
417            tmp = IB_SETUP.with_suffix(".tmp")
418            urllib.request.urlretrieve(
419                f"https://clients.bitrock.com/installbuilder/installbuilder-enterprise-{IB_VERSION}-windows-x64-installer.exe",
420                tmp,
421                reporthook=report
422            )
423            tmp.rename(IB_SETUP)
424
425        ib_setup_hash = hashlib.sha256()
426        with IB_SETUP.open("rb") as fp:
427            while True:
428                data = fp.read(65_536)
429                if not data:
430                    break
431                ib_setup_hash.update(data)
432        if ib_setup_hash.hexdigest() != IB_SETUP_SHA256:  # pragma: no cover
433            raise RuntimeError("InstallBuilder hashes don't match.")
434
435        click.echo("Install InstallBuilder...")
436        subprocess.run([IB_SETUP, "--mode", "unattended", "--unattendedmodeui", "none"], check=True)
437        assert IB_CLI.is_file()
438
439    if not IB_LICENSE.exists():
440        assert be.build_key
441        click.echo("Decrypt InstallBuilder license...")
442        f = cryptography.fernet.Fernet(be.build_key.encode())
443        with open(IB_LICENSE.with_suffix(".xml.enc"), "rb") as infile, \
444                open(IB_LICENSE, "wb") as outfile:
445            outfile.write(f.decrypt(infile.read()))
446
447    click.echo("Run InstallBuilder...")
448    subprocess.run([
449        IB_CLI,
450        "build",
451        str(IB_DIR / "mitmproxy.xml"),
452        "windows",
453        "--license", str(IB_LICENSE),
454        "--setvars", f"project.version={be.version}",
455        "--verbose"
456    ], check=True)
457    assert (be.dist_dir / f"mitmproxy-{be.version}-windows-installer.exe").exists()
458
459
460@click.group(chain=True)
461def cli():  # pragma: no cover
462    """
463    mitmproxy build tool
464    """
465    pass
466
467
468@cli.command("build")
469def build():  # pragma: no cover
470    """
471        Build a binary distribution
472    """
473    be = BuildEnviron.from_env()
474    be.dump_info()
475
476    be.check_version()
477    os.makedirs(be.dist_dir, exist_ok=True)
478
479    if be.should_build_wheel:
480        build_wheel(be)
481    if be.should_build_docker:
482        build_docker_image(be)
483    if be.should_build_pyinstaller:
484        build_pyinstaller(be)
485    if be.should_build_wininstaller:
486        build_wininstaller(be)
487
488
489@cli.command("upload")
490def upload():  # pragma: no cover
491    """
492        Upload build artifacts
493
494        Uploads the wheels package to PyPi.
495        Uploads the Pyinstaller and wheels packages to the snapshot server.
496        Pushes the Docker image to Docker Hub.
497    """
498    be = BuildEnviron.from_env()
499    be.dump_info()
500
501    if be.is_pull_request:
502        click.echo("Refusing to upload artifacts from a pull request!")
503        return
504
505    if be.should_upload_aws:
506        num_files = len([name for name in be.dist_dir.iterdir() if name.is_file()])
507        click.echo(f"Uploading {num_files} files to AWS dir {be.upload_dir}...")
508        subprocess.check_call([
509            "aws", "s3", "cp",
510            "--acl", "public-read",
511            f"{be.dist_dir}/",
512            f"s3://snapshots.mitmproxy.org/{be.upload_dir}/",
513            "--recursive",
514        ])
515
516    if be.should_upload_pypi:
517        whl, = be.dist_dir.glob('mitmproxy-*-py3-none-any.whl')
518        click.echo(f"Uploading {whl} to PyPi...")
519        subprocess.check_call(["twine", "upload", whl])
520
521    if be.should_upload_docker:
522        click.echo(f"Uploading Docker image to tag={be.docker_tag}...")
523        subprocess.check_call([
524            "docker",
525            "login",
526            "-u", be.docker_username,
527            "-p", be.docker_password,
528        ])
529
530        whl, = be.dist_dir.glob('mitmproxy-*-py3-none-any.whl')
531        docker_build_dir = be.release_dir / "docker"
532        shutil.copy(whl, docker_build_dir / whl.name)
533        # buildx is a bit weird in that we need to reinvoke build, but oh well.
534        subprocess.check_call([
535            "docker", "buildx", "build",
536            "--tag", be.docker_tag,
537            "--push",
538            "--platform", DOCKER_PLATFORMS,
539            "--build-arg", f"MITMPROXY_WHEEL={whl.name}",
540            "."
541        ], cwd=docker_build_dir)
542
543        if be.is_prod_release:
544            subprocess.check_call([
545                "docker", "buildx", "build",
546                "--tag", "mitmproxy/mitmproxy:latest",
547                "--push",
548                "--platform", DOCKER_PLATFORMS,
549                "--build-arg", f"MITMPROXY_WHEEL={whl.name}",
550                "."
551            ], cwd=docker_build_dir)
552
553
554if __name__ == "__main__":  # pragma: no cover
555    cli()
556