import tarfile
from os import walk

from empire.archive.abstract_archive import AbstractArchive
from empire.fs.paths import Paths
from empire.util.log import *


class TAR(AbstractArchive):
    @staticmethod
    def archive_directory(path: str, output_archive_file: str) -> bool:
        try:
            handle = tarfile.TarFile(path, 'w', format=tarfile.PAX_FORMAT)
        except Exception as error:
            Log.error('Unable to instantiate tarfile object', __file__, get_function_name(), error=error)
            return False

        TAR._tar_dir(path, handle)
        return True

    @staticmethod
    def unarchive_directory(archive_file: str, output_path: str) -> bool:
        """
        Warning!
        Never extract archives from untrusted sources without prior inspection.
        It is possible that files are created outside of path, e.g. members that have absolute
        filenames starting with "/" or filenames with two dots "..".
        """
        try:
            handle = tarfile.TarFile(archive_file, 'r')
        except Exception as error:
            Log.error('Unable to instantiate tarfile object', __file__, get_function_name(), error=error)
            return False

        try:
            handle.extractall(output_path)
            return True
        except Exception as error:
            Log.error('An error occurred while uncompressing "{}" to "{}"'.format(archive_file, output_path), __file__, get_function_name(), error=error)
            return False

    @staticmethod
    def archive_file(file: str, destination_archive_file: str) -> bool:
        try:
            handle = tarfile.TarFile(destination_archive_file, 'a')
        except Exception as error:
            Log.error('Unable to instantiate tarfile object', __file__, get_function_name(), error=error)
            return False

        try:
            handle.add(file)
            return True
        except Exception as error:
            Log.error('An error occurred while adding file "{}" to TAR archive'.format(file), __file__, get_function_name(), error=error)
            return False

    @staticmethod
    def unarchive_file(archive_file: str, archive_inner_file: str, destination_path: str) -> bool:
        try:
            handle = tarfile.TarFile(destination_path, 'r')
        except Exception as error:
            Log.error('Unable to instantiate tarfile object', __file__, get_function_name(), error=error)
            return False

        try:
            handle.extract(archive_inner_file, path=destination_path)
            return True
        except Exception as error:
            Log.error('An error occurred while extracting file "{}" from TAR archive "{}"'.format(archive_inner_file, archive_file), __file__, get_function_name(), error=error)
            return False

    @staticmethod
    def _tar_dir(path: str, tar_handle: tarfile.TarFile):
        for root, dirs, files in walk(path):
            for file in files:
                tar_handle.add(Paths.join(root, file), recursive=False)
