#!/usr/bin/env python
"""Test the grr aff4 objects."""

import StringIO
import time

import mock

from grr.lib import flags
from grr.lib import rdfvalue
from grr.lib import utils
from grr.lib.rdfvalues import client as rdf_client
from grr.lib.rdfvalues import cloud
from grr.lib.rdfvalues import flows as rdf_flows
from grr.lib.rdfvalues import paths as rdf_paths
from grr.server import aff4
from grr.server import data_store
from grr.server import flow
from grr.server.aff4_objects import aff4_grr
from grr.server.flows.general import transfer
from grr.test_lib import action_mocks
from grr.test_lib import aff4_test_lib
from grr.test_lib import flow_test_lib
from grr.test_lib import test_lib


class MockChangeEvent(flow.EventListener):
  EVENTS = ["MockChangeEvent"]

  well_known_session_id = rdfvalue.SessionID(flow_name="MockChangeEventHandler")

  CHANGED_URNS = []

  @flow.EventHandler(allow_client_access=True)
  def ProcessMessage(self, message=None, event=None):
    _ = event
    if (message.auth_state !=
        rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED):
      return

    urn = rdfvalue.RDFURN(message.payload)
    MockChangeEvent.CHANGED_URNS.append(urn)


class AFF4GRRTest(aff4_test_lib.AFF4ObjectTest):
  """Test the client aff4 implementation."""

  def setUp(self):
    super(AFF4GRRTest, self).setUp()
    MockChangeEvent.CHANGED_URNS = []

  def testAFF4Path(self):
    """Test the pathspec to URN conversion function."""
    pathspec = rdf_paths.PathSpec(
        path="\\\\.\\Volume{1234}\\",
        pathtype=rdf_paths.PathSpec.PathType.OS,
        mount_point="/c:/").Append(
            path="/windows", pathtype=rdf_paths.PathSpec.PathType.TSK)

    urn = pathspec.AFF4Path(rdf_client.ClientURN("C.1234567812345678"))
    self.assertEqual(
        urn,
        rdfvalue.RDFURN(
            r"aff4:/C.1234567812345678/fs/tsk/\\.\Volume{1234}\/windows"))

    # Test an ADS
    pathspec = rdf_paths.PathSpec(
        path="\\\\.\\Volume{1234}\\",
        pathtype=rdf_paths.PathSpec.PathType.OS,
        mount_point="/c:/").Append(
            pathtype=rdf_paths.PathSpec.PathType.TSK,
            path="/Test Directory/notes.txt:ads",
            inode=66,
            ntfs_type=128,
            ntfs_id=2)

    urn = pathspec.AFF4Path(rdf_client.ClientURN("C.1234567812345678"))
    self.assertEqual(urn,
                     rdfvalue.RDFURN(
                         r"aff4:/C.1234567812345678/fs/tsk/\\.\Volume{1234}\/"
                         "Test Directory/notes.txt:ads"))

  def testClientSubfieldGet(self):
    """Test we can get subfields of the client."""
    fd = aff4.FACTORY.Create(
        "C.0000000000000000", aff4_grr.VFSGRRClient, token=self.token)

    kb = fd.Schema.KNOWLEDGE_BASE()
    for i in range(5):
      kb.users.Append(rdf_client.User(username="user%s" % i))
    fd.Set(kb)
    fd.Close()

    fd = aff4.FACTORY.Open(
        "C.0000000000000000", aff4_grr.VFSGRRClient, token=self.token)
    for i, user in enumerate(fd.Get(fd.Schema.KNOWLEDGE_BASE).users):
      self.assertEqual(user.username, "user%s" % i)

  def testVFSFileContentLastNotUpdated(self):
    """Make sure CONTENT_LAST does not update when only STAT is written.."""
    path = "/C.12345/contentlastchecker"

    timestamp = 1
    with utils.Stubber(time, "time", lambda: timestamp):
      fd = aff4.FACTORY.Create(
          path, aff4_grr.VFSFile, mode="w", token=self.token)

      timestamp += 1
      fd.SetChunksize(10)

      # Make lots of small writes - The length of this string and the chunk size
      # are relative primes for worst case.
      for i in range(100):
        fd.Write("%s%08X\n" % ("Test", i))

        # Flush after every write.
        fd.Flush()

        # And advance the time.
        timestamp += 1

      fd.Set(fd.Schema.STAT, rdf_client.StatEntry())

      fd.Close()

    fd = aff4.FACTORY.Open(path, mode="rw", token=self.token)
    # Make sure the attribute was written when the write occured.
    self.assertEqual(int(fd.GetContentAge()), 101000000)

    # Write the stat (to be the same as before, but this still counts
    # as a write).
    fd.Set(fd.Schema.STAT, fd.Get(fd.Schema.STAT))
    fd.Flush()

    fd = aff4.FACTORY.Open(path, token=self.token)

    # The age of the content should still be the same.
    self.assertEqual(int(fd.GetContentAge()), 101000000)

  def testVFSFileStartsOnlyOneMultiGetFileFlowOnUpdate(self):
    """File updates should only start one MultiGetFile at any point in time."""
    client_id = self.SetupClient(0)
    # We need to create a file path having a pathspec.
    path = "fs/os/c/bin/bash"

    with aff4.FACTORY.Create(
        client_id.Add(path),
        aff4_type=aff4_grr.VFSFile,
        mode="rw",
        token=self.token) as file_fd:
      file_fd.Set(
          file_fd.Schema.STAT,
          rdf_client.StatEntry(
              pathspec=rdf_paths.PathSpec(path="/bin/bash", pathtype="OS")))

      # Starts a MultiGetFile flow.
      file_fd.Update()

    # Check that there is exactly one flow on the client.
    flows_fd = aff4.FACTORY.Open(client_id.Add("flows"), token=self.token)
    flows = list(flows_fd.ListChildren())
    self.assertEqual(len(flows), 1)

    # The flow is the MultiGetFile flow holding the lock on the file.
    flow_obj = aff4.FACTORY.Open(flows[0], token=self.token)
    self.assertEqual(
        flow_obj.Get(flow_obj.Schema.TYPE), transfer.MultiGetFile.__name__)
    self.assertEqual(flow_obj.urn, file_fd.Get(file_fd.Schema.CONTENT_LOCK))

    # Since there is already a running flow having the lock on the file,
    # this call shouldn't do anything.
    file_fd.Update()

    # There should still be only one flow on the client.
    flows_fd = aff4.FACTORY.Open(client_id.Add("flows"), token=self.token)
    flows = list(flows_fd.ListChildren())
    self.assertEqual(len(flows), 1)

  def testVFSFileStartsNewMultiGetFileWhenLockingFlowHasFinished(self):
    """A new MultiFileGet can be started when the locking flow has finished."""
    client_id = self.SetupClient(0)
    path = "fs/os/c/bin/bash"

    with aff4.FACTORY.Create(
        client_id.Add(path),
        aff4_type=aff4_grr.VFSFile,
        mode="rw",
        token=self.token) as file_fd:
      file_fd.Set(
          file_fd.Schema.STAT,
          rdf_client.StatEntry(
              pathspec=rdf_paths.PathSpec(path="/bin/bash", pathtype="OS")))
      # Starts a MultiGetFile flow.
      first_update_flow_urn = file_fd.Update()

    # Check that there is exactly one flow on the client.
    flows_fd = aff4.FACTORY.Open(client_id.Add("flows"), token=self.token)
    flows = list(flows_fd.ListChildren())
    self.assertEqual(len(flows), 1)

    # Finish the flow holding the lock.
    client_mock = action_mocks.ActionMock()
    for _ in flow_test_lib.TestFlowHelper(
        flows[0], client_mock, client_id=client_id, token=self.token):
      pass

    # The flow holding the lock has finished, so Update() should start a new
    # flow.
    second_update_flow_urn = file_fd.Update()

    # There should be two flows now.
    flows_fd = aff4.FACTORY.Open(client_id.Add("flows"), token=self.token)
    flows = list(flows_fd.ListChildren())
    self.assertEqual(len(flows), 2)

    # Make sure that each Update() started a new flow and that the second flow
    # is holding the lock.
    self.assertNotEqual(first_update_flow_urn, second_update_flow_urn)
    self.assertEqual(second_update_flow_urn,
                     file_fd.Get(file_fd.Schema.CONTENT_LOCK))

  def testGetClientSummary(self):
    hostname = "test"
    system = "Linux"
    os_release = "12.02"
    kernel = "3.15-rc2"
    fqdn = "test.test.com"
    arch = "amd64"
    install_time = rdfvalue.RDFDatetime.Now()
    user = "testuser"
    userobj = rdf_client.User(username=user)
    interface = rdf_client.Interface(ifname="eth0")
    google_cloud_instance = cloud.GoogleCloudInstance(
        instance_id="1771384456894610289",
        zone="projects/123456789733/zones/us-central1-a",
        project_id="myproject",
        unique_id="us-central1-a/myproject/1771384456894610289")
    cloud_instance = cloud.CloudInstance(
        cloud_type="GOOGLE", google=google_cloud_instance)

    serial_number = "DSD33679FZ"
    system_manufacturer = "Foobar Inc."
    system_uuid = "C31292AD-6Z4F-55D8-28AC-EC1100E42222"
    hwinfo = rdf_client.HardwareInfo(
        serial_number=serial_number,
        system_manufacturer=system_manufacturer,
        system_uuid=system_uuid)

    timestamp = 1
    with utils.Stubber(time, "time", lambda: timestamp):
      with aff4.FACTORY.Create(
          "C.0000000000000000",
          aff4_grr.VFSGRRClient,
          mode="rw",
          token=self.token) as fd:
        kb = rdf_client.KnowledgeBase()
        kb.users.Append(userobj)
        empty_summary = fd.GetSummary()
        self.assertEqual(empty_summary.client_id, "C.0000000000000000")
        self.assertFalse(empty_summary.system_info.version)
        self.assertEqual(empty_summary.timestamp.AsSecondsFromEpoch(), 1)

        # This will cause TYPE to be written with current time = 101 when the
        # object is closed
        timestamp += 100
        fd.Set(fd.Schema.HOSTNAME(hostname))
        fd.Set(fd.Schema.SYSTEM(system))
        fd.Set(fd.Schema.OS_RELEASE(os_release))
        fd.Set(fd.Schema.KERNEL(kernel))
        fd.Set(fd.Schema.FQDN(fqdn))
        fd.Set(fd.Schema.ARCH(arch))
        fd.Set(fd.Schema.INSTALL_DATE(install_time))
        fd.Set(fd.Schema.KNOWLEDGE_BASE(kb))
        fd.Set(fd.Schema.USERNAMES([user]))
        fd.Set(fd.Schema.HARDWARE_INFO(hwinfo))
        fd.Set(fd.Schema.INTERFACES([interface]))
        fd.Set(fd.Schema.CLOUD_INSTANCE(cloud_instance))

      with aff4.FACTORY.Open(
          "C.0000000000000000",
          aff4_grr.VFSGRRClient,
          mode="rw",
          token=self.token) as fd:
        summary = fd.GetSummary()
        self.assertEqual(summary.system_info.system, system)
        self.assertEqual(summary.system_info.release, os_release)
        self.assertEqual(summary.system_info.kernel, kernel)
        self.assertEqual(summary.system_info.fqdn, fqdn)
        self.assertEqual(summary.system_info.machine, arch)
        self.assertEqual(summary.system_info.install_date, install_time)
        self.assertItemsEqual(summary.users, [userobj])
        self.assertItemsEqual(summary.interfaces, [interface])
        self.assertFalse(summary.client_info)

        self.assertEqual(summary.timestamp.AsSecondsFromEpoch(), 101)
        self.assertEqual(summary.cloud_type, "GOOGLE")
        self.assertEqual(summary.cloud_instance_id,
                         "us-central1-a/myproject/1771384456894610289")

        self.assertEqual(summary.serial_number, serial_number)
        self.assertEqual(summary.system_manufacturer, system_manufacturer)
        self.assertEqual(summary.system_uuid, system_uuid)


class BlobImageTest(aff4_test_lib.AFF4ObjectTest):
  """Tests for cron functionality."""

  def testAppendContentError(self):
    src_content = "ABCD" * 10
    src_fd = StringIO.StringIO(src_content)

    dest_fd = aff4.FACTORY.Create(
        aff4.ROOT_URN.Add("temp"),
        aff4_grr.VFSBlobImage,
        token=self.token,
        mode="rw")
    dest_fd.SetChunksize(7)
    dest_fd.AppendContent(src_fd)
    dest_fd.Seek(0)
    self.assertEqual(dest_fd.Read(5000), src_content)

    src_fd.seek(0)
    self.assertRaises(IOError, dest_fd.AppendContent, src_fd)

  def testAppendContent(self):
    """Test writing content where content length % chunksize == 0."""
    src_content = "ABCDEFG" * 10  # 10 chunksize blobs
    src_fd = StringIO.StringIO(src_content)

    dest_fd = aff4.FACTORY.Create(
        aff4.ROOT_URN.Add("temp"),
        aff4_grr.VFSBlobImage,
        token=self.token,
        mode="rw")
    self.assertEqual(dest_fd.Get(dest_fd.Schema.HASHES), None)

    dest_fd.SetChunksize(7)
    dest_fd.AppendContent(src_fd)

    self.assertEqual(int(dest_fd.Get(dest_fd.Schema.SIZE)), len(src_content))
    self.assertTrue(dest_fd.Get(dest_fd.Schema.HASHES))

    dest_fd.Seek(0)
    self.assertEqual(dest_fd.Read(5000), src_content)

    src_fd.seek(0)
    dest_fd.AppendContent(src_fd)
    self.assertEqual(dest_fd.size, 2 * len(src_content))
    self.assertEqual(
        int(dest_fd.Get(dest_fd.Schema.SIZE)), 2 * len(src_content))
    dest_fd.Seek(0)
    self.assertEqual(dest_fd.Read(5000), src_content + src_content)

  def testMultiStreamStreamsSingleFileWithSingleChunk(self):
    with aff4.FACTORY.Create(
        "aff4:/foo", aff4_type=aff4_grr.VFSBlobImage, token=self.token) as fd:
      fd.SetChunksize(10)
      fd.AppendContent(StringIO.StringIO("123456789"))

    fd = aff4.FACTORY.Open("aff4:/foo", token=self.token)
    chunks_fds = list(aff4.AFF4Stream.MultiStream([fd]))

    self.assertEqual(len(chunks_fds), 1)
    self.assertEqual(chunks_fds[0][1], "123456789")
    self.assertIs(chunks_fds[0][0], fd)

  def testMultiStreamStreamsSinglfeFileWithTwoChunks(self):
    with aff4.FACTORY.Create(
        "aff4:/foo", aff4_type=aff4_grr.VFSBlobImage, token=self.token) as fd:
      fd.SetChunksize(10)
      fd.AppendContent(StringIO.StringIO("123456789"))

    with aff4.FACTORY.Create(
        "aff4:/bar", aff4_type=aff4_grr.VFSBlobImage, token=self.token) as fd:
      fd.SetChunksize(10)
      fd.AppendContent(StringIO.StringIO("abcd"))

    fd1 = aff4.FACTORY.Open("aff4:/foo", token=self.token)
    fd2 = aff4.FACTORY.Open("aff4:/bar", token=self.token)
    chunks_fds = list(aff4.AFF4Stream.MultiStream([fd1, fd2]))

    self.assertEqual(len(chunks_fds), 2)

    self.assertEqual(chunks_fds[0][1], "123456789")
    self.assertIs(chunks_fds[0][0], fd1)

    self.assertEqual(chunks_fds[1][1], "abcd")
    self.assertIs(chunks_fds[1][0], fd2)

  def testMultiStreamStreamsTwoFilesWithTwoChunksInEach(self):
    with aff4.FACTORY.Create(
        "aff4:/foo", aff4_type=aff4_grr.VFSBlobImage, token=self.token) as fd:
      fd.SetChunksize(10)
      fd.AppendContent(StringIO.StringIO("*" * 10 + "123456789"))

    with aff4.FACTORY.Create(
        "aff4:/bar", aff4_type=aff4_grr.VFSBlobImage, token=self.token) as fd:
      fd.SetChunksize(10)
      fd.AppendContent(StringIO.StringIO("*" * 10 + "abcd"))

    fd1 = aff4.FACTORY.Open("aff4:/foo", token=self.token)
    fd2 = aff4.FACTORY.Open("aff4:/bar", token=self.token)
    chunks_fds = list(aff4.AFF4Stream.MultiStream([fd1, fd2]))

    self.assertEqual(len(chunks_fds), 4)

    self.assertEqual(chunks_fds[0][1], "*" * 10)
    self.assertIs(chunks_fds[0][0], fd1)

    self.assertEqual(chunks_fds[1][1], "123456789")
    self.assertIs(chunks_fds[1][0], fd1)

    self.assertEqual(chunks_fds[2][1], "*" * 10)
    self.assertIs(chunks_fds[2][0], fd2)

    self.assertEqual(chunks_fds[3][1], "abcd")
    self.assertIs(chunks_fds[3][0], fd2)

  def testMultiStreamReturnsExceptionIfChunkIsMissing(self):
    with aff4.FACTORY.Create(
        "aff4:/foo", aff4_type=aff4_grr.VFSBlobImage, token=self.token) as fd:
      fd.SetChunksize(10)
      fd.AppendContent(StringIO.StringIO("123456789"))

      fd.index.seek(0)
      blob_id = fd.index.read(fd._HASH_SIZE).encode("hex")

    data_store.DB.DeleteBlob(blob_id, token=self.token)

    fd = aff4.FACTORY.Open("aff4:/foo", token=self.token)
    returned_fd, _, e = list(aff4.AFF4Stream.MultiStream([fd]))[0]
    self.assertNotEqual(e, None)
    self.assertEqual(returned_fd, fd)
    self.assertEqual(e.missing_chunks, [blob_id])

  def testMultiStreamIgnoresTheFileIfAnyChunkIsMissingInReadAheadChunks(self):
    with aff4.FACTORY.Create(
        "aff4:/foo", aff4_type=aff4_grr.VFSBlobImage, token=self.token) as fd:
      fd.SetChunksize(10)
      fd.AppendContent(StringIO.StringIO("*" * 10 + "123456789"))

      fd.index.seek(0)
      unused_blob_id_1 = fd.index.read(fd._HASH_SIZE).encode("hex")
      blob_id_2 = fd.index.read(fd._HASH_SIZE).encode("hex")

    data_store.DB.DeleteBlob(blob_id_2, token=self.token)

    fd = aff4.FACTORY.Open("aff4:/foo", token=self.token)
    count = 0
    for _, _, e in aff4.AFF4Stream.MultiStream([fd]):
      if not e:
        count += 1

    self.assertEqual(count, 0)

  @mock.patch.object(aff4_grr.VFSBlobImage, "MULTI_STREAM_CHUNKS_READ_AHEAD", 1)
  def testMultiStreamTruncatesBigFileIfLastChunkIsMissing(self):
    # If the file is split between 2 batches of chunks, and the missing
    # chunk is in the second batch, the first batch will be succesfully
    # yielded.
    with aff4.FACTORY.Create(
        "aff4:/foo", aff4_type=aff4_grr.VFSBlobImage, token=self.token) as fd:
      fd.SetChunksize(10)
      fd.AppendContent(StringIO.StringIO("*" * 10 + "123456789"))

      fd.index.seek(0)
      unused_blob_id_1 = fd.index.read(fd._HASH_SIZE).encode("hex")
      blob_id_2 = fd.index.read(fd._HASH_SIZE).encode("hex")

    data_store.DB.DeleteBlob(blob_id_2, token=self.token)

    fd = aff4.FACTORY.Open("aff4:/foo", token=self.token)
    content = []
    error_detected = False
    for fd, chunk, e in aff4.AFF4Stream.MultiStream([fd]):
      if not e:
        content.append(chunk)
      else:
        error_detected = True

    self.assertEqual(content, ["*" * 10])
    self.assertTrue(error_detected)

  @mock.patch.object(aff4_grr.VFSBlobImage, "MULTI_STREAM_CHUNKS_READ_AHEAD", 1)
  def testMultiStreamSkipsBigFileIfFirstChunkIsMissing(self):
    # If the file is split between 2 batches of chunks, and the missing
    # chunk is in the first batch, the file will be skipped entirely.
    with aff4.FACTORY.Create(
        "aff4:/foo", aff4_type=aff4_grr.VFSBlobImage, token=self.token) as fd:
      fd.SetChunksize(10)
      fd.AppendContent(StringIO.StringIO("*" * 10 + "123456789"))

      fd.index.seek(0)
      blob_id_1 = fd.index.read(fd._HASH_SIZE).encode("hex")

    data_store.DB.DeleteBlob(blob_id_1, token=self.token)

    fd = aff4.FACTORY.Open("aff4:/foo", token=self.token)
    count = 0
    for _, _, e in aff4.AFF4Stream.MultiStream([fd]):
      if not e:
        count += 1

    self.assertEqual(count, 0)


def main(argv):
  # Run the full test suite
  test_lib.main(argv)


if __name__ == "__main__":
  flags.StartMain(main)
