# beipack - Remote bootloader for Python
#
# Copyright (C) 2022 Allison Karlitskaya <allison.karlitskaya@redhat.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
import argparse
import binascii
import lzma
import os
import sys
import tempfile
import zipfile
from typing import Dict, Iterable, List, Optional, Set, Tuple
from .data import read_data_file
def escape_string(data: str) -> str:
# Avoid mentioning ' ' ' literally, to make our own packing a bit prettier
triplequote = "'" * 3
if triplequote not in data:
return "r" + triplequote + data + triplequote
if '"""' not in data:
return 'r"""' + data + '"""'
return repr(data)
def ascii_bytes_repr(data: bytes) -> str:
return 'b' + escape_string(data.decode('ascii'))
def utf8_bytes_repr(data: bytes) -> str:
return escape_string(data.decode('utf-8')) + ".encode('utf-8')"
def base64_bytes_repr(data: bytes, imports: Set[str]) -> str:
# base85 is smaller, but base64 is in C, and ~20x faster.
# when compressing with `xz -e` the size difference is marginal.
imports.add('from binascii import a2b_base64')
encoded = binascii.b2a_base64(data).decode('ascii').strip()
return f'a2b_base64("{encoded}")'
def bytes_repr(data: bytes, imports: Set[str]) -> str:
# Strategy:
# if the file is ascii, encode it directly as bytes
# otherwise, if it's UTF-8, use a unicode string and encode
# otherwise, base64
try:
return ascii_bytes_repr(data)
except UnicodeDecodeError:
# it's not ascii
pass
# utf-8
try:
return utf8_bytes_repr(data)
except UnicodeDecodeError:
# it's not utf-8
pass
return base64_bytes_repr(data, imports)
def dict_repr(contents: Dict[str, bytes], imports: Set[str]) -> str:
return ('{\n' +
''.join(f' {repr(k)}: {bytes_repr(v, imports)},\n'
for k, v in contents.items()) +
'}')
def pack(contents: Dict[str, bytes],
entrypoint: Optional[str] = None,
args: str = '') -> str:
"""Creates a beipack with the given `contents`.
If `entrypoint` is given, it should be an entry point which is run as the
"main" function. It is given in the `package.module:func format` such that
the following code is emitted:
from package.module import func as main
main()
Additionally, if `args` is given, it is written verbatim between the parens
of the call to main (ie: it should already be in Python syntax).
"""
loader = read_data_file('beipack_loader.py')
lines = [line for line in loader.splitlines() if line]
lines.append('')
imports = {'import sys'}
contents_txt = dict_repr(contents, imports)
lines.extend(imports)
lines.append(f'sys.meta_path.insert(0, BeipackLoader({contents_txt}))')
if entrypoint:
package, main = entrypoint.split(':')
lines.append(f'from {package} import {main} as main')
lines.append(f'main({args})')
return ''.join(f'{line}\n' for line in lines)
def collect_contents(filenames: List[str],
relative_to: Optional[str] = None) -> Dict[str, bytes]:
contents: Dict[str, bytes] = {}
for filename in filenames:
with open(filename, 'rb') as file:
contents[os.path.relpath(filename, start=relative_to)] = file.read()
return contents
def collect_module(name: str, *, recursive: bool) -> Dict[str, bytes]:
import importlib.resources
from importlib.resources.abc import Traversable
def walk(path: str, entry: Traversable) -> Iterable[Tuple[str, bytes]]:
for item in entry.iterdir():
itemname = f'{path}/{item.name}'
if item.is_file():
yield itemname, item.read_bytes()
elif recursive and item.name != '__pycache__':
yield from walk(itemname, item)
return dict(walk(name.replace('.', '/'), importlib.resources.files(name)))
def collect_zip(filename: str) -> Dict[str, bytes]:
contents = {}
with zipfile.ZipFile(filename) as file:
for entry in file.filelist:
if '.dist-info/' in entry.filename:
continue
contents[entry.filename] = file.read(entry)
return contents
def collect_pep517(path: str) -> Dict[str, bytes]:
with tempfile.TemporaryDirectory() as tmpdir:
import build
builder = build.ProjectBuilder(path)
wheel = builder.build('wheel', tmpdir)
return collect_zip(wheel)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument('--python', '-p',
help="add a #!python3 interpreter line using the given path")
parser.add_argument('--xz', '-J', action='store_true',
help="compress the output with `xz`")
parser.add_argument('--topdir',
help="toplevel directory (paths are stored relative to this)")
parser.add_argument('--output', '-o',
help="write output to a file (default: stdout)")
parser.add_argument('--main', '-m', metavar='MODULE:FUNC',
help="use FUNC from MODULE as the main function")
parser.add_argument('--main-args', metavar='ARGS',
help="arguments to main() in Python syntax", default='')
parser.add_argument('--module', action='append', default=[],
help="collect installed modules (recursively)")
parser.add_argument('--zip', '-z', action='append', default=[],
help="include files from a zipfile (or wheel)")
parser.add_argument('--build', metavar='DIR', action='append', default=[],
help="PEP-517 from a given source directory")
parser.add_argument('files', nargs='*',
help="files to include in the beipack")
args = parser.parse_args()
contents = collect_contents(args.files, relative_to=args.topdir)
for file in args.zip:
contents.update(collect_zip(file))
for name in args.module:
contents.update(collect_module(name, recursive=True))
for path in args.build:
contents.update(collect_pep517(path))
result = pack(contents, args.main, args.main_args).encode('utf-8')
if args.python:
result = b'#!' + args.python.encode('ascii') + b'\n' + result
if args.xz:
result = lzma.compress(result, preset=lzma.PRESET_EXTREME)
if args.output:
with open(args.output, 'wb') as file:
file.write(result)
else:
if args.xz and os.isatty(1):
sys.exit('refusing to write compressed output to a terminal')
sys.stdout.buffer.write(result)
if __name__ == '__main__':
main()