import argparse
import lzma
import os
import subprocess
import tempfile
from typing import List, Dict, Optional

import requests
import rpmfile
from debian import debfile

DWARF2JSON = "./dwarf2json"


class Downloader:

    def __init__(self, url_lists: List[List[str]]) -> None:
        self.url_lists = url_lists

    def download_lists(self, keep=False):
        for url_list in self.url_lists:
            print("Downloading files...")
            files_for_processing = self.download_list(url_list)
            self.process_files(files_for_processing)
            if not keep:
                for fname in files_for_processing.values():
                    if fname:
                        os.unlink(fname)

    def download_list(self, urls: List[str]) -> Dict[str, str]:
        processed_files = {}
        for url in urls:
            print(f" - Downloading {url}")
            data = requests.get(url)
            with tempfile.NamedTemporaryFile() as archivedata:
                archivedata.write(data.content)
                archivedata.seek(0)
                if url.endswith(".rpm"):
                    processed_files[url] = self.process_rpm(archivedata)
                elif url.endswith(".deb"):
                    processed_files[url] = self.process_deb(archivedata)

        return processed_files

    def process_rpm(self, archivedata) -> Optional[str]:
        rpm = rpmfile.RPMFile(fileobj=archivedata)
        member = None
        extracted = None
        for member in rpm.getmembers():
            if "vmlinux" in member.name or "System.map" in member.name:
                print(f" - Extracting {member.name}")
                extracted = rpm.extractfile(member)
                break
        if not member or not extracted:
            return None
        with tempfile.NamedTemporaryFile(
            delete=False, prefix="vmlinux" if "vmlinux" in member.name else "System.map"
        ) as output:
            print(f" - Writing to {output.name}")
            output.write(extracted.read())
        return output.name

    def process_deb(self, archivedata) -> Optional[str]:
        deb = debfile.DebFile(fileobj=archivedata)
        member = None
        extracted = None
        for member in deb.data.tgz().getmembers():
            if member.name.endswith("vmlinux") or "System.map" in member.name:
                print(f" - Extracting {member.name}")
                extracted = deb.data.get_file(member.name)
                break
        if not member or not extracted:
            return None
        with tempfile.NamedTemporaryFile(
            delete=False, prefix="vmlinux" if "vmlinux" in member.name else "System.map"
        ) as output:
            print(f" - Writing to {output.name}")
            output.write(extracted.read())
        return output.name

    def process_files(self, named_files: Dict[str, str]):
        """Runs the dwarf2json binary across the files"""
        print("Processing Files...")
        for i in named_files:
            if named_files[i] is None:
                print(f"FAILURE: None encountered for {i}")
                return
        args = [DWARF2JSON, "linux"]
        output_filename = "unknown-kernel.json"
        for named_file in named_files:
            prefix = "--system-map"
            if "System" not in named_files[named_file]:
                prefix = "--elf"
                output_filename = (
                    "./"
                    + "-".join((named_file.split("/")[-1]).split("-")[2:])[:-4]
                    + ".json.xz"
                )
            args += [prefix, named_files[named_file]]
        print(f" - Running {args}")
        proc = subprocess.run(args, capture_output=True)

        print(f" - Writing to {output_filename}")
        with lzma.open(output_filename, "w") as f:
            f.write(proc.stdout)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Takes a list of URLs for Centos and downloads them"
    )
    parser.add_argument(
        "-f",
        "--file",
        dest="filename",
        metavar="FILENAME",
        help="Filename to be read",
        required=True,
    )
    parser.add_argument(
        "-d",
        "--dwarf2json",
        dest="dwarfpath",
        metavar="PATH",
        default=DWARF2JSON,
        help="Path to the dwarf2json binary",
        required=True,
    )
    parser.add_argument(
        "-k",
        "--keep",
        dest="keep",
        action="store_true",
        help="Keep extracted temporary files after completion",
        default=False,
    )
    args = parser.parse_args()

    DWARF2JSON = args.dwarfpath

    with open(args.filename) as f:
        lines = f.readlines()

    urls = []
    for i in range(len(lines) // 2):
        urls += [[lines[2 * i].strip(), lines[(2 * i) + 1].strip()]]

    d = Downloader(urls)
    d.download_lists(keep=args.keep)