#!/usr/bin/env python3

import argparse

import base64

import hashlib

import os.path

import shutil

import tempfile

import zipfile

from typing import NamedTuple

class Wheel(NamedTuple):

src: str

plat: str

exe: str = 'sentry-cli'

WHEELS = (

Wheel(

src='sentry-cli-Darwin-arm64',

plat='macosx_11_0_arm64',

),

Wheel(

src='sentry-cli-Darwin-universal',

plat='macosx_11_0_universal2',

),

Wheel(

src='sentry-cli-Darwin-x86_64',

plat='macosx_10_15_x86_64',

),

Wheel(

src='sentry-cli-Linux-aarch64',

plat='manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_2_aarch64',

),

Wheel(

src='sentry-cli-Linux-armv7',

plat='manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_2_armv7l',

),

Wheel(

src='sentry-cli-Linux-i686',

plat='manylinux_2_17_i686.manylinux2014_i686.musllinux_1_2_i686',

),

Wheel(

src='sentry-cli-Linux-x86_64',

plat='manylinux_2_17_x86_64.manylinux2014_x86_64.musllinux_1_2_x86_64',

),

Wheel(

src='sentry-cli-Windows-i686.exe',

plat='win32',

exe='sentry-cli.exe',

),

Wheel(

src='sentry-cli-Windows-x86_64.exe',

plat='win_amd64',

exe='sentry-cli.exe',

),

Wheel(

src='sentry-cli-Windows-aarch64.exe',

plat='win_arm64',

exe='sentry-cli.exe',

),

)

def main() -> int:

parser = argparse.ArgumentParser()

parser.add_argument('--binaries', required=True)

parser.add_argument('--base', required=True)

parser.add_argument('--dest', required=True)

args = parser.parse_args()

expected = {wheel.src for wheel in WHEELS}

received = set(os.listdir(args.binaries))

if expected < received:

raise SystemExit(

f'Unexpected binaries:\n\n'

f'- extra: {", ".join(sorted(received - expected))}\n'

f'- missing: {", ".join(sorted(expected - received))}'

)

sdist_path = wheel_path = None

for fname in os.listdir(args.base):

if fname.endswith('.tar.gz'):

sdist_path = os.path.join(args.base, fname)

elif fname.endswith('.whl'):

wheel_path = os.path.join(args.base, fname)

else:

raise SystemExit(f'unexpected file in `--base`: {fname}')

if sdist_path is None or wheel_path is None:

raise SystemExit('expected wheel and sdist in `--base`')

os.makedirs(args.dest, exist_ok=True)

shutil.copy(sdist_path, args.dest)

for wheel in WHEELS:

binary_src = os.path.join(args.binaries, wheel.src)

binary_size = os.stat(binary_src).st_size

with open(binary_src, 'rb') as bf:

digest = hashlib.sha256(bf.read()).digest()

digest_b64 = base64.urlsafe_b64encode(digest).rstrip(b'=').decode()

basename = os.path.basename(wheel_path)

wheelname, _ = os.path.splitext(basename)

name, version, py, abi, plat = wheelname.split('-')

with tempfile.TemporaryDirectory() as tmp:

with zipfile.ZipFile(wheel_path) as zipf:

zipf.extractall(tmp)

distinfo = os.path.join(tmp, f'{name}-{version}.dist-info')

scripts = os.path.join(tmp, f'{name}-{version}.data', 'scripts')

# replace the script binary with our copy

os.remove(os.path.join(scripts, 'sentry-cli'))

shutil.copy(binary_src, os.path.join(scripts, wheel.exe))

# rewrite RECORD to include the new file

record_fname = os.path.join(distinfo, 'RECORD')

with open(record_fname) as f:

record_lines = list(f)

record = f'{name}-{version}.data/scripts/sentry-cli,'

for i, line in enumerate(record_lines):

if line.startswith(record):

record_lines[i] = (

f'{name}-{version}.data/scripts/{wheel.exe},'

f'sha256={digest_b64},'

f'{binary_size}\n'

)

break

else:

raise SystemExit(f'could not find {record!r} in RECORD')

with open(record_fname, 'w') as f:

f.writelines(record_lines)

# rewrite WHEEL to have the new tags

wheel_fname = os.path.join(distinfo, 'WHEEL')

with open(wheel_fname) as f:

wheel_lines = list(f)

for i, line in enumerate(wheel_lines):

if line.startswith('Tag: '):

wheel_lines[i:i + 1] = [

f'Tag: {py}-{abi}-{plat}\n'

for plat in wheel.plat.split('.')

]

break

else:

raise SystemExit("could not find 'Tag: ' in WHEEL")

with open(wheel_fname, 'w') as f:

f.writelines(wheel_lines)

# write out the final zip

new_basename = f'{name}-{version}-{py}-{abi}-{wheel.plat}.whl'

tmp_new_wheel = os.path.join(tmp, new_basename)

fnames = sorted(

os.path.join(root, fname)

for root, _, fnames in os.walk(tmp)

for fname in fnames

)

with zipfile.ZipFile(tmp_new_wheel, 'w') as zipf:

for fname in fnames:

zinfo = zipfile.ZipInfo(os.path.relpath(fname, tmp))

if '/scripts/' in zinfo.filename:

zinfo.external_attr = 0o100755 << 16

with open(fname, 'rb') as fb:

zipf.writestr(zinfo, fb.read())

# move into dest

shutil.move(tmp_new_wheel, args.dest)

return 0

if __name__ == '__main__':

raise SystemExit(main())