#!/usr/bin/env python
"""Contains tests for archive_generator."""

import hashlib
import os
import tarfile
from typing import Iterable
from unittest import mock
import zipfile

from absl import app
import yaml

from grr_response_proto import flows_pb2
from grr_response_proto import jobs_pb2
from grr_response_proto import objects_pb2
from grr_response_server import data_store
from grr_response_server import file_store
from grr_response_server import flow_base
from grr_response_server.databases import db
from grr_response_server.gui import archive_generator
from grr_response_server.models import blobs
from grr_response_server.rdfvalues import mig_flow_objects
from grr_response_server.rdfvalues import objects as rdf_objects
from grr.test_lib import flow_test_lib
from grr.test_lib import test_lib
from grr.test_lib import vfs_test_lib


class CollectionArchiveGeneratorTest(test_lib.GRRBaseTest):
  """Test for CollectionArchiveGenerator."""

  def setUp(self):
    super().setUp()
    self.client_id = self.SetupClient(0)

  def _CreateFile(self, client_id: str, vfs_path: str, content: bytes) -> None:
    digest = hashlib.sha256(content).digest()
    path_type, components = rdf_objects.ParseCategorizedPath(vfs_path)

    path_info = objects_pb2.PathInfo(
        path_type=path_type,
        components=components,
    )

    blob_id = blobs.BlobID.Of(digest)
    data_store.BLOBS.WriteBlobs({blob_id: content})
    blob_ref = rdf_objects.BlobReference(
        offset=0, size=len(content), blob_id=bytes(blob_id)
    )
    hash_id = file_store.AddFileWithUnknownHash(
        db.ClientPath.FromPathInfo(client_id, path_info), [blob_ref]
    )
    path_info.hash_entry.sha256 = hash_id.AsBytes()

    data_store.REL_DB.WritePathInfos(client_id, [path_info])

  def _InitializeFiles(self) -> None:
    path1 = "fs/os/foo/bar/hello1.txt"
    archive_path1 = "test_prefix/%s/fs/os/foo/bar/hello1.txt" % self.client_id
    self._CreateFile(
        client_id=self.client_id,
        vfs_path=path1,
        content="hello1".encode("utf-8"),
    )

    path2 = "fs/os/foo/bar/中国新闻网新闻中.txt"
    archive_path2 = (
        "test_prefix/%s/fs/os/foo/bar/中国新闻网新闻中.txt" % self.client_id
    )
    self._CreateFile(
        client_id=self.client_id,
        vfs_path=path2,
        content="hello2".encode("utf-8"),
    )

    self.flow_results = []
    self.paths = [path1, path2]
    self.archive_paths = [archive_path1, archive_path2]
    for path in self.paths:
      flow_result = flows_pb2.FlowResult()
      flow_result.client_id = self.client_id
      flow_result.payload.Pack(
          jobs_pb2.StatEntry(
              pathspec=jobs_pb2.PathSpec(
                  path="foo/bar/" + str(path).split("/")[-1],
                  pathtype=jobs_pb2.PathSpec.PathType.OS,
              )
          )
      )
      self.flow_results.append(flow_result)

  def _GenerateArchive(
      self,
      collection: Iterable[flows_pb2.FlowResult],
      archive_format=archive_generator.ArchiveFormat.ZIP,
      predicate=None,
  ):

    fd_path = os.path.join(self.temp_dir, "archive")
    generator = archive_generator.CollectionArchiveGenerator(
        archive_format=archive_format,
        predicate=predicate,
        prefix="test_prefix",
        description="Test description",
    )
    with open(fd_path, "wb") as out_fd:
      for chunk in generator.Generate(collection):
        out_fd.write(chunk)

    return fd_path

  def testCreatesZipContainingFilesAndClientInfosAndManifest(self):
    self._InitializeFiles()

    fd_path = self._GenerateArchive(
        self.flow_results,
        archive_format=archive_generator.ArchiveFormat.ZIP,
    )

    zip_fd = zipfile.ZipFile(fd_path)
    names = [str(s) for s in sorted(zip_fd.namelist())]

    client_info_name = "test_prefix/%s/client_info.yaml" % self.client_id
    manifest_name = "test_prefix/MANIFEST"

    self.assertCountEqual(
        names, self.archive_paths + [client_info_name, manifest_name]
    )

    contents = zip_fd.read(self.archive_paths[0])
    self.assertEqual(contents, b"hello1")

    contents = zip_fd.read(self.archive_paths[1])
    self.assertEqual(contents, b"hello2")

    manifest = yaml.safe_load(zip_fd.read(manifest_name))
    self.assertEqual(
        manifest,
        {
            "description": "Test description",
            "processed_files": 2,
            "archived_files": 2,
            "ignored_files": 0,
            "failed_files": 0,
        },
    )

    client_info = yaml.safe_load(zip_fd.read(client_info_name))

    self.assertEqual(
        client_info["knowledge_base"]["fqdn"], "Host-0.example.com"
    )

  def testCreatesTarContainingFilesAndClientInfosAndManifest(self):
    self._InitializeFiles()

    fd_path = self._GenerateArchive(
        self.flow_results,
        archive_format=archive_generator.ArchiveFormat.TAR_GZ,
    )

    with tarfile.open(fd_path, encoding="utf-8") as tar_fd:
      manifest_fd = tar_fd.extractfile("test_prefix/MANIFEST")
      self.assertEqual(
          yaml.safe_load(manifest_fd.read()),
          {
              "description": "Test description",
              "processed_files": 2,
              "archived_files": 2,
              "ignored_files": 0,
              "failed_files": 0,
          },
      )

      archive_path_0, archive_path_1 = self.archive_paths

      self.assertEqual(tar_fd.extractfile(archive_path_0).read(), b"hello1")
      self.assertEqual(tar_fd.extractfile(archive_path_1).read(), b"hello2")

      client_info_name = "test_prefix/%s/client_info.yaml" % self.client_id
      client_info = yaml.safe_load(tar_fd.extractfile(client_info_name).read())

      self.assertEqual(
          client_info["knowledge_base"]["fqdn"], "Host-0.example.com"
      )

  def testCorrectlyAccountsForFailedFiles(self):
    self._InitializeFiles()

    with mock.patch.object(
        file_store, "StreamFilesChunks", side_effect=Exception("foobar")
    ):
      with self.assertRaises(Exception) as context:
        self._GenerateArchive(
            self.flow_results,
            archive_format=archive_generator.ArchiveFormat.ZIP,
        )
      self.assertEqual(str(context.exception), "foobar")

  def testNotFoundFilesProduceWarning(self):
    self._InitializeFiles()
    flow_results = list(self.flow_results)

    flow_result = flows_pb2.FlowResult()
    flow_result.client_id = self.client_id
    flow_result.payload.Pack(
        jobs_pb2.StatEntry(
            pathspec=jobs_pb2.PathSpec(
                path="foo/bar/notfound", pathtype=jobs_pb2.PathSpec.PathType.OS
            )
        )
    )
    flow_results.append(flow_result)

    fd_path = self._GenerateArchive(
        flow_results,
        archive_format=archive_generator.ArchiveFormat.ZIP,
    )

    zip_fd = zipfile.ZipFile(fd_path)
    yaml_str = zip_fd.read("test_prefix/MANIFEST")
    manifest = yaml.safe_load(yaml_str)
    self.assertEqual(
        manifest,
        {
            "description": "Test description",
            "processed_files": 3,
            "archived_files": 2,
            "ignored_files": 0,
            "failed_files": 0,
        },
    )
    self.assertIn(
        archive_generator.CollectionArchiveGenerator.FILES_SKIPPED_WARNING,
        yaml_str,
    )

  def testIgnoresFilesNotMatchingPredicate(self):
    self._InitializeFiles()

    def predicate(pathspec):
      return os.path.basename(pathspec.Path()).startswith("hello")

    fd_path = self._GenerateArchive(
        self.flow_results,
        predicate=predicate,
        archive_format=archive_generator.ArchiveFormat.ZIP,
    )

    zip_fd = zipfile.ZipFile(fd_path)
    names = sorted(zip_fd.namelist())

    # The archive is expected to contain 1 file contents blob, 1 client info and
    # a manifest.
    self.assertLen(names, 3)

    manifest = yaml.safe_load(zip_fd.read("test_prefix/MANIFEST"))
    self.assertEqual(
        manifest,
        {
            "description": "Test description",
            "processed_files": 2,
            "archived_files": 1,
            "ignored_files": 1,
            "failed_files": 0,
            "ignored_files_list": [
                "aff4:/%s/fs/os/foo/bar/中国新闻网新闻中.txt" % self.client_id
            ],
        },
    )


class FlowArchiveGeneratorTest(test_lib.GRRBaseTest):
  """Test for CollectionArchiveGenerator."""

  def _GenerateArchive(self, generator: Iterable[bytes]) -> str:
    fd_path = os.path.join(self.temp_dir, "archive")
    with open(fd_path, "wb") as out_fd:
      for chunk in generator:
        out_fd.write(chunk)

    return fd_path

  def setUp(self):
    super().setUp()

    self.client_id = self.SetupClient(0)
    self.flow_id = flow_test_lib.StartFlow(
        flow_test_lib.DummyFlow, client_id=self.client_id
    )
    self.flow = data_store.REL_DB.ReadFlowObject(self.client_id, self.flow_id)
    self.flow = mig_flow_objects.ToRDFFlow(self.flow)

    self.path1 = db.ClientPath.OS(self.client_id, ["foo", "bar", "hello1.txt"])
    self.path1_content = "hello1".encode("utf-8")
    self.path2 = db.ClientPath.TSK(
        self.client_id, ["foo", "bar", "中国新闻网新闻中.txt"]
    )
    self.path2_content = "hello2".encode("utf-8")

    vfs_test_lib.CreateFile(self.path1, self.path1_content)
    vfs_test_lib.CreateFile(self.path2, self.path2_content)

  def testCreatesZipContainingTwoMappedFilesAndManifest(self):
    generator = archive_generator.FlowArchiveGenerator(
        self.flow, archive_generator.ArchiveFormat.ZIP
    )
    mappings = [
        flow_base.ClientPathArchiveMapping(self.path1, "foo/file"),
        flow_base.ClientPathArchiveMapping(self.path2, "foo/bar/file"),
    ]
    fd_path = self._GenerateArchive(generator.Generate(mappings))

    zip_fd = zipfile.ZipFile(fd_path)
    names = [str(s) for s in sorted(zip_fd.namelist())]

    # Expecting in the archive: 2 files and a manifest.
    self.assertLen(names, 3)

    contents = zip_fd.read(os.path.join(generator.prefix, "foo", "file"))
    self.assertEqual(contents, b"hello1")

    contents = zip_fd.read(os.path.join(generator.prefix, "foo", "bar", "file"))
    self.assertEqual(contents, b"hello2")

    manifest = yaml.safe_load(
        zip_fd.read(os.path.join(generator.prefix, "MANIFEST"))
    )
    self.assertCountEqual(
        manifest["processed_files"].items(),
        [
            (self.path1.vfs_path, "foo/file"),
            (self.path2.vfs_path, "foo/bar/file"),
        ],
    )
    self.assertCountEqual(manifest["missing_files"], [])
    self.assertEqual(manifest["client_id"], self.client_id)
    self.assertEqual(manifest["flow_id"], self.flow_id)

  def testCreatesTarContainingTwoMappedFilesAndManifest(self):
    generator = archive_generator.FlowArchiveGenerator(
        self.flow, archive_generator.ArchiveFormat.TAR_GZ
    )
    mappings = [
        flow_base.ClientPathArchiveMapping(self.path1, "foo/file"),
        flow_base.ClientPathArchiveMapping(self.path2, "foo/bar/file"),
    ]
    fd_path = self._GenerateArchive(generator.Generate(mappings))

    with tarfile.open(fd_path, encoding="utf-8") as tar_fd:
      self.assertLen(tar_fd.getnames(), 3)

      contents = tar_fd.extractfile(
          os.path.join(generator.prefix, "foo", "file")
      ).read()
      self.assertEqual(contents, b"hello1")

      contents = tar_fd.extractfile(
          os.path.join(generator.prefix, "foo", "bar", "file")
      ).read()
      self.assertEqual(contents, b"hello2")

      manifest = yaml.safe_load(
          tar_fd.extractfile(os.path.join(generator.prefix, "MANIFEST")).read()
      )
      self.assertCountEqual(
          manifest["processed_files"].items(),
          [
              (self.path1.vfs_path, "foo/file"),
              (self.path2.vfs_path, "foo/bar/file"),
          ],
      )
      self.assertCountEqual(manifest["missing_files"], [])
      self.assertEqual(manifest["client_id"], self.client_id)
      self.assertEqual(manifest["flow_id"], self.flow_id)

  def testPropagatesStreamingExceptions(self):
    generator = archive_generator.FlowArchiveGenerator(
        self.flow, archive_generator.ArchiveFormat.TAR_GZ
    )
    mappings = [
        flow_base.ClientPathArchiveMapping(self.path1, "foo/file"),
        flow_base.ClientPathArchiveMapping(self.path2, "foo/bar/file"),
    ]

    with mock.patch.object(
        file_store, "StreamFilesChunks", side_effect=Exception("foobar")
    ):
      with self.assertRaises(Exception) as context:
        self._GenerateArchive(generator.Generate(mappings))
      self.assertEqual(str(context.exception), "foobar")

  def testMissingFilesAreListedInManifest(self):
    generator = archive_generator.FlowArchiveGenerator(
        self.flow, archive_generator.ArchiveFormat.ZIP
    )
    mappings = [
        flow_base.ClientPathArchiveMapping(self.path1, "foo/file"),
        flow_base.ClientPathArchiveMapping(
            db.ClientPath.OS(self.client_id, ["non", "existing"]),
            "foo/bar/file",
        ),
    ]
    fd_path = self._GenerateArchive(generator.Generate(mappings))

    zip_fd = zipfile.ZipFile(fd_path)
    names = [str(s) for s in sorted(zip_fd.namelist())]

    # Expecting in the archive: 1 file (the other shouldn't be found)
    # and a manifest.
    self.assertLen(names, 2)

    contents = zip_fd.read(os.path.join(generator.prefix, "foo", "file"))
    self.assertEqual(contents, b"hello1")

    manifest = yaml.safe_load(
        zip_fd.read(os.path.join(generator.prefix, "MANIFEST"))
    )
    self.assertCountEqual(
        manifest["processed_files"].items(), [(self.path1.vfs_path, "foo/file")]
    )
    self.assertCountEqual(manifest["missing_files"], ["fs/os/non/existing"])


def main(argv):
  test_lib.main(argv)


if __name__ == "__main__":
  app.run(main)
