# Copyright 2023 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Helper functions for dealing with .zip files."""

import os
import pathlib
import posixpath
import stat
import time
import zipfile

_FIXED_ZIP_HEADER_LEN = 30


def _set_alignment(zip_obj, zip_info, alignment):
    """Sets a ZipInfo's extra field such that the file will be aligned.

    Args:
      zip_obj: The ZipFile object that is being written.
      zip_info: The ZipInfo object about to be written.
      alignment: The amount of alignment (e.g. 4, or 4*1024).
    """
    header_size = _FIXED_ZIP_HEADER_LEN + len(zip_info.filename)
    pos = zip_obj.fp.tell() + header_size
    padding_needed = (alignment - (pos % alignment)) % alignment

    # Python writes |extra| to both the local file header and the central
    # directory's file header. Android's zipalign tool writes only to the
    # local file header, so there is more overhead in using Python to align.
    zip_info.extra = b'\0' * padding_needed


def _hermetic_date_time(timestamp=None):
    if not timestamp:
        return (2001, 1, 1, 0, 0, 0)
    utc_time = time.gmtime(timestamp)
    return (
        utc_time.tm_year,
        utc_time.tm_mon,
        utc_time.tm_mday,
        utc_time.tm_hour,
        utc_time.tm_min,
        utc_time.tm_sec,
    )


def add_to_zip_hermetic(
    zip_file,
    zip_path,
    *,
    src_path=None,
    data=None,
    compress=None,
    compress_level=1,
    alignment=None,
    timestamp=None,
):
    """Adds a file to the given ZipFile with a hard-coded modified time.

    Args:
      zip_file: ZipFile instance to add the file to.
      zip_path: Destination path within the zip file (or ZipInfo instance).
      src_path: Path of the source file. Mutually exclusive with |data|.
      data: File data as a string.
      compress: Whether to enable compression. Default is taken from ZipFile
          constructor.
      compress_level: When compress=True, level of compression.
      alignment: If set, align the data of the entry to this many bytes.
      timestamp: The last modification date and time for the archive member.
    """
    assert (src_path is None) != (data is None), (
        '|src_path| and |data| are mutually exclusive.'
    )
    if isinstance(zip_path, zipfile.ZipInfo):
        zipinfo = zip_path
        zip_path = zipinfo.filename
    else:
        zipinfo = zipfile.ZipInfo(filename=zip_path)
        zipinfo.external_attr = 0o644 << 16

    zipinfo.date_time = _hermetic_date_time(timestamp)

    if alignment:
        _set_alignment(zip_file, zipinfo, alignment)

    # Filenames can contain backslashes, but it is more likely that we've
    # forgotten to use forward slashes as a directory separator.
    assert '\\' not in zip_path, 'zip_path should not contain \\: ' + zip_path
    assert not posixpath.isabs(zip_path), 'Absolute zip path: ' + zip_path
    assert not zip_path.startswith('..'), (
        'Should not start with ..: ' + zip_path
    )
    assert posixpath.normpath(zip_path) == zip_path, (
        f'Non-canonical zip_path: {zip_path} vs: {posixpath.normpath(zip_path)}'
    )
    assert zip_path not in zip_file.namelist(), (
        'Tried to add a duplicate zip entry: ' + zip_path
    )

    if src_path and os.path.islink(src_path):
        zipinfo.external_attr |= stat.S_IFLNK << 16  # mark as a symlink
        zip_file.writestr(zipinfo, os.readlink(src_path))
        return

    # Maintain the executable bit.
    if src_path:
        st = os.stat(src_path)
        for mode in (stat.S_IXUSR, stat.S_IXGRP, stat.S_IXOTH):
            if st.st_mode & mode:
                zipinfo.external_attr |= mode << 16

    if src_path:
        with open(src_path, 'rb') as f:
            data = f.read()

    # zipfile will deflate even when it makes the file bigger. To avoid
    # growing files, disable compression at an arbitrary cut off point.
    if len(data) < 16:
        compress = False

    # None converts to ZIP_STORED, when passed explicitly rather than the
    # default passed to the ZipFile constructor.
    compress_type = zip_file.compression
    if compress is not None:
        compress_type = zipfile.ZIP_DEFLATED if compress else zipfile.ZIP_STORED
    zip_file.writestr(
        zipinfo, data, compress_type, compresslevel=compress_level
    )


def add_files_to_zip(
    inputs,
    output,
    *,
    base_dir=None,
    path_transform=None,
    compress=None,
    compress_level=1,
    zip_prefix_path=None,
    timestamp=None,
):
    """Creates a zip file from a list of files.

    Args:
      inputs: A list of paths to zip, or a list of (zip_path, fs_path) tuples.
      output: Path, fileobj, or ZipFile instance to add files to.
      base_dir: Prefix to strip from inputs.
      path_transform: Called for each entry path. Returns a new zip path, or None
          to skip the file.
      compress: Whether to compress
      compress_level: When compress=True, level of compression.
      zip_prefix_path: Path prepended to file path in zip file.
      timestamp: Unix timestamp to use for files in the archive.
    """
    if base_dir is None:
        base_dir = '.'
    input_tuples = []
    for tup in inputs:
        if isinstance(tup, str):
            src_path = tup
            zip_path = os.path.relpath(src_path, base_dir)
            # Zip files always use / as path separator.
            if os.path.sep != posixpath.sep:
                zip_path = str(pathlib.Path(zip_path).as_posix())
            tup = (zip_path, src_path)
        input_tuples.append(tup)

    # Sort by zip path to ensure stable zip ordering.
    input_tuples.sort(key=lambda tup: tup[0])

    out_zip = output
    if not isinstance(output, zipfile.ZipFile):
        out_zip = zipfile.ZipFile(output, 'w')

    try:
        for zip_path, fs_path in input_tuples:
            if zip_prefix_path:
                zip_path = posixpath.join(zip_prefix_path, zip_path)
            if path_transform:
                zip_path = path_transform(zip_path)
                if zip_path is None:
                    continue
            add_to_zip_hermetic(
                out_zip,
                zip_path,
                src_path=fs_path,
                compress=compress,
                compress_level=compress_level,
                timestamp=timestamp,
            )
    finally:
        if output is not out_zip:
            out_zip.close()


def zip_directory(output, base_dir, **kwargs):
    """Zips all files in the given directory."""
    inputs = []
    for root, _, files in os.walk(base_dir):
        for f in files:
            inputs.append(os.path.join(root, f))

    add_files_to_zip(inputs, output, base_dir=base_dir, **kwargs)


def merge_zips(
    output, input_zips, path_transform=None, compress=None, compress_level=1
):
    """Combines all files from |input_zips| into |output|.

    Args:
      output: Path, fileobj, or ZipFile instance to add files to.
      input_zips: Iterable of paths to zip files to merge.
      path_transform: Called for each entry path. Returns a new zip path, or None
          to skip the file.
      compress: Overrides compression setting from origin zip entries.
      compress_level: When compress=True, level of compression.
    """
    assert not isinstance(input_zips, str)  # Easy mistake to make.
    if isinstance(output, zipfile.ZipFile):
        out_zip = output
        out_filename = output.filename
    else:
        assert isinstance(output, str), 'Was: ' + repr(output)
        out_zip = zipfile.ZipFile(output, 'w')
        out_filename = output

    # Include paths in the existing zip here to avoid adding duplicate files.
    crc_by_name = {
        i.filename: (out_filename, i.CRC) for i in out_zip.infolist()
    }

    try:
        for in_file in input_zips:
            with zipfile.ZipFile(in_file, 'r') as in_zip:
                for info in in_zip.infolist():
                    # Ignore directories.
                    if info.filename[-1] == '/':
                        continue
                    if path_transform:
                        dst_name = path_transform(info.filename)
                        if dst_name is None:
                            continue
                    else:
                        dst_name = info.filename

                    data = in_zip.read(info)

                    # If there's a duplicate file, ensure contents is the same and skip
                    # adding it multiple times.
                    if dst_name in crc_by_name:
                        orig_filename, orig_crc = crc_by_name[dst_name]
                        new_crc = zipfile.crc32(data)
                        if new_crc == orig_crc:
                            continue
                        msg = f"""File appeared in multiple inputs with differing contents.
File: {dst_name}
Input1: {orig_filename}
Input2: {in_file}"""
                        raise Exception(msg)

                    if compress is not None:
                        compress_entry = compress
                    else:
                        compress_entry = (
                            info.compress_type != zipfile.ZIP_STORED
                        )
                    add_to_zip_hermetic(
                        out_zip,
                        dst_name,
                        data=data,
                        compress=compress_entry,
                        compress_level=compress_level,
                    )
                    crc_by_name[dst_name] = (
                        in_file,
                        out_zip.getinfo(dst_name).CRC,
                    )
    finally:
        if output is not out_zip:
            out_zip.close()