import argparse
import base64
import json
import logging
import os
import pathlib
import urllib

from volatility3.cli import PrintedProgress
from volatility3.framework import contexts, constants
from volatility3.framework.automagic import linux, mac

vollog = logging.getLogger(__name__)


class BannerCacheGenerator:

    def __init__(self, path: str, url_prefix: str):
        self._path = path
        self._url_prefix = url_prefix

    def convert_url(self, url):
        parsed = urllib.parse.urlparse(url)

        relpath = os.path.relpath(parsed.path, os.path.abspath(self._path))

        return urllib.parse.urljoin(self._url_prefix, relpath)

    def run(self):
        context = contexts.Context()
        json_output = {'version': 1}

        path = self._path
        filename = '*'

        for banner_cache in [linux.LinuxBannerCache, mac.MacBannerCache]:
            sub_path = banner_cache.os
            potentials = []
            for extension in constants.ISF_EXTENSIONS:
                # Hopefully these will not be large lists, otherwise this might be slow
                try:
                    for found in pathlib.Path(path).joinpath(sub_path).resolve().rglob(filename + extension):
                        potentials.append(found.as_uri())
                except FileNotFoundError:
                    # If there's no linux symbols, don't cry about it
                    pass

            new_banners = banner_cache.read_new_banners(context, 'BannerServer', potentials, banner_cache.symbol_name,
                                                        banner_cache.os, progress_callback = PrintedProgress())
            result_banners = {}
            for new_banner in new_banners:
                # Only accept file schemes
                value = [self.convert_url(url) for url in new_banners[new_banner] if
                         urllib.parse.urlparse(url).scheme == 'file']
                if value and new_banner:
                    # Convert files into URLs
                    result_banners[str(base64.b64encode(new_banner), 'latin-1')] = value

            json_output[banner_cache.os] = result_banners

        output_path = os.path.join(self._path, 'banners.json')
        with open(output_path, 'w') as fp:
            vollog.warning(f"Banners file written to {output_path}")
            json.dump(json_output, fp)


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--path', default = os.path.dirname(__file__))
    parser.add_argument('--urlprefix', help = 'Web prefix that will eventually serve the ISF files',
                        default = 'http://localhost/symbols')

    args = parser.parse_args()

    bcg = BannerCacheGenerator(args.path, args.urlprefix)
    bcg.run()