#!/usr/bin/env python3
# Copyright 2018 Ciena Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Download and extract a Kafka distribution tarball.
To list the supported Kafka versions:

    tools/download-kafka list

To download a given version:

    tools/download-kafka get <KAFKA_VERSION>
"""

import argparse
import hashlib
import os
import sys
import tarfile
import time
from shutil import copyfileobj
from urllib.request import urlopen

CHUNK_SIZE = 32 * 1024 * 1024

# This table lists the Kafka versions this script knows how to download. Each
# entry is a 3-tuple of (archive root directory, download URL, and sha512
# hash).
versions = {
    '0.8.0': (
        'kafka_2.8.0-0.8.0',
        'https://archive.apache.org/dist/kafka/0.8.0/kafka_2.8.0-0.8.0.tar.gz',
        '4eab03a98384415d8fd7ec6c5708a25492d05342ee7bdfc7ff4fccc715b79ba0dd248fef797272d7a3a9ae1a002a00d1d8ea50b4546f38956411b3cc5b3a6c3e',
    ),
    '0.8.1': (
        'kafka_2.9.2-0.8.1',
        'https://archive.apache.org/dist/kafka/0.8.1/kafka_2.9.2-0.8.1.tgz',
        '75dc932126b9c7a6bf005ccfa39c22d09c9d0f395a2fcc682814922a254e539cf6c8b869266cefc047dbc0a991d340d9fa212cd5c272de6fbe500d1287cc55d8',
    ),
    '0.8.1.1': (
        'kafka_2.9.2-0.8.1.1',
        'https://archive.apache.org/dist/kafka/0.8.1.1/kafka_2.9.2-0.8.1.1.tgz',
        '9897e899ecb3ddbd71705436ed7e568166685b5d168f2c2f852956aacfb80a93b6864c1502c9a4537fc27b587e1fa06e2390c0f4f56d7400fef116245c4f1d81',
    ),
    '0.8.2.1': (
        'kafka_2.10-0.8.2.1',
        'https://archive.apache.org/dist/kafka/0.8.2.1/kafka_2.10-0.8.2.1.tgz',
        '48d73a18eae615a5aeb52b1158fd9a4bb96b7a3c50fe52599ffb88586f10c422c2f0078958cafe1c63518970f1e3984fc70c95a3b2896514d9d02eb7e119aef8',
    ),
    '0.8.2.2': (
        'kafka_2.10-0.8.2.2',
        'https://archive.apache.org/dist/kafka/0.8.2.2/kafka_2.10-0.8.2.2.tgz',
        '03978c4e82257eb36cc4c80c82c8e22a18a4546fd857b4fa34b089674960894321efb9c461229944b2b148c32c0cc4ec8a5d4c01ea8aee67ed9bedcf632a6b0d',
    ),
    '0.9.0.1': (
        'kafka_2.11-0.9.0.1',
        'https://archive.apache.org/dist/kafka/0.9.0.1/kafka_2.11-0.9.0.1.tgz',
        '818a821877436bbb9698a32afeb299bc569985a6b712179a9882bf1a10edfdb4992519ba39c1d7499144650b30781fc943a4d5d15b8d5916a4c1a5e542d9b956',
    ),
    '0.11.0.3': (
        'kafka_2.12-0.11.0.3',
        'https://archive.apache.org/dist/kafka/0.11.0.3/kafka_2.11-0.11.0.3.tgz',
        'd38caaa80f43d02dcc8bc453fbf71e8d609249731583556fdd991dcb09ff342d0ec855896ff76875cea48a471cc95bda9174bf3f3507696f243e72e5e456c584',
    ),
    '1.0.0': (
        'kafka_2.11-1.0.0',
        'https://archive.apache.org/dist/kafka/1.0.0/kafka_2.11-1.0.0.tgz',
        'b698028485d460c494af942b2dfa41c7342a2ad6052b543c95289cd1e832a1ebcf679b7e568dcabc342c799765337b94f91bc0cf0ee915534cdf82d1635a622a',
    ),
    '1.1.0': (
        'kafka_2.11-1.1.0',
        'https://archive.apache.org/dist/kafka/1.1.0/kafka_2.11-1.1.0.tgz',
        '5e613dbb21c2aed797e18bb5638386267c4f61ee00d10760dbc6bd472af8b6aa9cdca30c1f93f5e31a3ad8c506374a7350357e2ac2b9a3ddf87caf22820d470e',
    ),
    '1.1.1': (
        'kafka_2.11-1.1.1',
        'https://archive.apache.org/dist/kafka/1.1.1/kafka_2.11-1.1.1.tgz',
        '2bf14a07c569c84736271471a9abb7b937b311780ed2a3d969ac0123737319e9151e0a69d6b8bd309a57b92cb00a90400e8e19e0512a6ee9206b2c91826af930',
    ),
    '2.8.2': (
        'kafka_2.12-2.8.2',
        'https://archive.apache.org/dist/kafka/2.8.2/kafka_2.12-2.8.2.tgz',
        'ce90f4722d49d896643ce63e1b4f8f1620c7a3dca603646733052b3a9155f3e3f7e93422015c0c588988be836487f9d638472a624f1c4895662daa3aa3ba2bb7',
    ),
    '3.1.0': (
        'kafka_2.12-3.1.0',
        'https://archive.apache.org/dist/kafka/3.1.0/kafka_2.12-3.1.0.tgz',
        'bdc05a784c4b5bbcfdfd3f552727a0b1de2a01c794b22495de04fd478f704d27e9751f5e77eca1cea8ebca22cd51d859029dc4858e8858f977142b7cd948f5e7'
    ),
    '3.4.1': (
        'kafka_2.12-3.4.1',
        'https://archive.apache.org/dist/kafka/3.4.1/kafka_2.12-3.4.1.tgz',
        'dfc0daf75874a1ceda24f8a03d71073e423596e4a912aea886374d7131be9f2c8788c1f76dd652e121e5325b72a920bb1886c399da412111f5f388c991de8406'
    )
}

parser = argparse.ArgumentParser(description="Download and extract Kafka distribution tarballs.")
parser.add_argument('--list', action='store_true', help="List all know Kafka versions", default=False)
parser.add_argument('--all', action='store_true', help="Download all known Kafka verisons", default=False)
parser.add_argument('version', nargs='*', help="Kafka version to download (unless --all is specified)", default=[])
args = parser.parse_args()

if args.list:
    print('\n'.join(versions.keys()))
    sys.exit(0)
elif args.all:
    download_versions = list(versions.items())
else:
    download_versions = [(v, versions[v]) for v in args.version]


def check_hash(filepath, digest):
    """
    Verify that the given file matches the given digest.

    :raises ValueError: on failed hash verification
    """
    hasher = hashlib.sha512()
    with open(filepath, 'rb') as f:
        while True:
            chunk = f.read(CHUNK_SIZE)
            if not chunk:
                break
            hasher.update(chunk)
    if hasher.digest() != bytes.fromhex(digest):
        raise ValueError('Digest {} of {} does not match expected value {}'.format(
            hasher.digest().hex(), filepath, digest))


server_dir = os.path.normpath(os.path.join(__file__, '..', '..', 'servers'))
cache_dir = os.path.join(server_dir, 'dist')
for version, (archive_root, url, digest) in download_versions:
    dist_name = url.rsplit('/', 1)[1]
    dist_path = os.path.join(cache_dir, dist_name)
    extract_path = os.path.join(server_dir, version)
    bin_path = os.path.join(extract_path, 'kafka-bin')

    if os.path.isdir(bin_path):
        print("Kafka", version, "is available in", bin_path, "(remove this tree to force re-extraction)")
        continue

    failures = 0
    while True:
        try:
            check_hash(dist_path, digest)
        except FileNotFoundError:
            pass  # No cached file.
        except ValueError:
            # The file failed hash verification.
            failures += 1
            if failures >= 3:
                raise Exception("Failed to download Kafka {}.".format(version, dist_path))
            else:
                print("Kafka", version, "download", dist_path, "is corrupt; deleting it.")
                os.unlink(dist_path)
        else:
            break

        print('Downloading Kafka', version, 'to', dist_path)
        with open(dist_path, 'wb') as fout:
            with urlopen(url, timeout=30) as fin:
                if fin.status != 200:
                    print('Request failed: HTTP', fin.status, fin.reason, fin.geturl())
                    print('Will retry in 10 seconds')
                    time.sleep(10)
                    continue
                copyfileobj(fin, fout)

    # And also extract it.
    print('Extracting Kafka', version, 'to', extract_path)
    with tarfile.open(dist_path) as archive:
        archive.extractall(extract_path)
    # Then rename it to the well-known "kafka-bin" location.
    os.rename(os.path.join(extract_path, archive_root), bin_path)
