#!/usr/bin/env python
"""Manage file uploads."""
import collections
import gzip
import hashlib
import StringIO
import struct
import sys
import zlib


from grr.lib.rdfvalues import crypto
from grr.lib.rdfvalues import flows
from grr.lib.rdfvalues import structs as rdf_structs

from grr.proto import jobs_pb2

# The header struct format.
HEADER_MAGIC = "GRR\x01"
PART_TYPE_ENCRYPTED_CIPHER = 1
PART_TYPE_ENCRYPTED_DATA = 2
HEADER_FMT = "4sIHIIIII"
HEADER_SIZE = struct.calcsize(HEADER_FMT)
EncryptedHeader = collections.namedtuple(
    "EncryptedHeader",
    [
        "Magic",  # HEADER_MAGIC
        "Version",  # Header Version.
        "PartType",  # One of PART_TYPE_*
        "PartLength",  # Length of entire part including header.
        "HMACStart",  # Offset of where the HMAC begins.
        "HMACLength",  # Length of HMAC.
        "DataStart",  # Offset where the part data begins.
        "DataLength",  # Length of data.
    ])


class SignaturePart(rdf_structs.RDFProtoStruct):
  """A Signature part."""
  protobuf = jobs_pb2.SignaturePart


class EncryptStream(object):
  """A file like object which encrypts and signs another stream."""

  def __init__(self,
               readers_public_key,
               writers_private_key,
               fd,
               chunk_size=1024 * 1024):
    """Constructor.

    Args:

      readers_public_key: The public key of the destined reader of this
        stream. The stream will be encrypted with the readers public key.
      writers_private_key: The private_key of the writer of this stream. Data
        will be signed using this private key.
      fd: A file like object we read from.
      chunk_size: This will be the size of the parts.
    """
    self.fd = fd
    self.readers_public_key = readers_public_key
    self.writers_private_key = writers_private_key
    self.chunk_size = chunk_size

    # Prepare the initial header.
    self.cipher_properties = flows.CipherProperties.GetInializedKeys()
    self.cipher = crypto.AES128CBCCipher(self.cipher_properties.key,
                                         self.cipher_properties.metadata_iv)
    self.hmac = crypto.HMAC(self.cipher_properties.hmac_key.RawBytes())

    serialized_cipher = self.cipher_properties.SerializeToString()
    signature = SignaturePart(
        encrypted_cipher=readers_public_key.Encrypt(serialized_cipher),
        signature=writers_private_key.Sign(serialized_cipher))

    # First part is the encrypted cipher.
    self.encrypted_buffer = BufferedReader()
    self.encrypted_buffer.write(
        self._GetPart(signature.SerializeToString(),
                      PART_TYPE_ENCRYPTED_CIPHER))
    self.eof = False

  def _GetPart(self, data, data_type):
    """Create a new part."""
    # Calculate a HMAC over the entire part.
    hmac = self.hmac.HMAC(data)
    header = EncryptedHeader(
        Magic=HEADER_MAGIC,
        Version=1,
        PartType=data_type,
        DataStart=HEADER_SIZE,
        DataLength=len(data),
        HMACStart=HEADER_SIZE + len(data),
        HMACLength=len(hmac),
        PartLength=HEADER_SIZE + len(hmac) + len(data))
    return struct.pack(HEADER_FMT, *header) + data + hmac

  def read(self, length):  # pylint: disable=invalid-name
    """Read length bytes from the stream."""
    if length is None:
      raise RuntimeError("Read calls must be bounded.")

    if self.eof:
      return ""

    while self.encrypted_buffer.len < length:
      current_chunk = self._EncryptNextChunk()
      # None indicates that the input stream is exhausted.
      if current_chunk is None:
        self.eof = True
        break

      self.encrypted_buffer.write(current_chunk)

    return self.encrypted_buffer.ReadFromFront(length)

  def Read(self, length):
    return self.read(length)

  def _EncryptNextChunk(self):
    """Encrypt a new chunk to send."""
    data = self.fd.read(self.chunk_size)
    if not data:
      return None

    # Send an encrypted data part.
    encrypted_data = self.cipher.Encrypt(data)
    return self._GetPart(encrypted_data, PART_TYPE_ENCRYPTED_DATA)

  def __enter__(self):
    return self

  def __exit__(self, exc_type, exc_value, traceback):
    return self.close()


class DecryptStream(object):
  """A file like object which decrypts another stream.

  This wrapper is designed to be inserted in the write() path of another
  stream. It will collect all parts written into it and in turn write those to
  the wrapped stream.

  """

  def __init__(self, readers_private_key, writers_public_key, outfd):
    """Constructor.

    Args:
      readers_private_key: The private key of the destined reader of this
        stream. The stream will be decrypted with the readers private key.
      writers_public_key: The public_key of the writer of this stream. Data's
        signature will be verified using this public key.
      outfd: A file like object we write to. This must be opened for writing.
    """
    self.outfd = outfd
    self.readers_private_key = readers_private_key
    self.writers_public_key = writers_public_key
    self.header = None
    self.part_end = None
    self.buffer = BufferedReader()
    self.cipher = None

  def Write(self, data):
    """Write some data into the stream."""
    self.buffer.write(data)

    while 1:
      # We are looking for the next part header.
      if self.part_end is None:
        # Not enough data to process a part.
        if self.buffer.len < HEADER_SIZE:
          return

        # Parse the header.
        packed_header = self.buffer.getvalue()[:HEADER_SIZE]
        self.header = EncryptedHeader(*struct.unpack(HEADER_FMT, packed_header))
        if self.header.Magic != HEADER_MAGIC or self.header.Version != 1:
          raise IOError("Invalid header.")

        # We need to wait until we have this much data in the buffer.
        self.part_end = self.header.PartLength

      # We do not have a full part yet.
      if self.buffer.len < self.part_end:
        return

      self._ProcessPart()

  def _ProcessPart(self):
    """Process a full part in the buffer.

    We assume a complete part exists in self.buffer.

    Raises:
      IOError: if there is anything wrong with the data.
    """
    payload_data = self.buffer.getvalue()[
        self.header.DataStart:self.header.DataStart + self.header.DataLength]

    payload_hmac = self.buffer.getvalue()[
        self.header.HMACStart:self.header.HMACStart + self.header.HMACLength]

    # Now process the part.
    # First header - we need to initialize the cipher.
    if self.cipher is None:
      if self.header.PartType != PART_TYPE_ENCRYPTED_CIPHER:
        raise IOError("First part must be ENCRYPTED_CIPHER")

      signature = SignaturePart.FromSerializedString(payload_data)
      decrypted_cipher = self.readers_private_key.Decrypt(
          signature.encrypted_cipher)

      # Verify its signature.
      self.writers_public_key.Verify(decrypted_cipher, signature.signature)

      self.cipher_properties = flows.CipherProperties.FromSerializedString(
          decrypted_cipher)

      self.cipher = self.cipher_properties.GetCipher()
      self.hmac = self.cipher_properties.GetHMAC()

    # The part contains data.
    elif self.header.PartType == PART_TYPE_ENCRYPTED_DATA:
      plain_text = self.cipher.Decrypt(payload_data)
      self.outfd.write(plain_text)

    else:
      raise IOError("Unsupported part type %s" % self.header.PartType)

    try:
      # Make sure the hmac is correct.
      self.hmac.Verify(payload_data, payload_hmac)
    except crypto.VerificationError:
      raise IOError("HMAC not verified")

    # Remove the part from the buffer.
    self.buffer.ReadFromFront(self.part_end)
    self.part_end = None

  def Close(self):
    if self.buffer.len > 0:
      raise IOError("Partial Message Received")
    self.flush()
    self.outfd.close()

  def Flush(self):
    self.outfd.flush()

  close = Close
  flush = Flush
  write = Write

  def __enter__(self):
    return self

  def __exit__(self, exc_type, exc_value, traceback):
    return self.close()


class BufferedReader(StringIO.StringIO):
  """A FLO which can be written to the back and read from the front."""

  def ReadFromFront(self, length):
    result_data = self.getvalue()
    self.truncate(0)
    self.write(result_data[length:])
    return result_data[:length]


class GzipWrapper(object):
  """Wraps a file like object and compresses it."""
  BUFFER_SIZE = 1024 * 1024

  def __init__(self, infd, byte_limit=None):
    self.total_read = 0
    self.infd = infd
    self.buff_fd = BufferedReader()
    self.zipper = gzip.GzipFile(mode="wb", fileobj=self.buff_fd)
    self.remaining_bytes = byte_limit
    if self.remaining_bytes is None:
      self.remaining_bytes = sys.maxint

    self.hashers = {
        "sha256": hashlib.sha256(),
        "sha1": hashlib.sha1(),
        "md5": hashlib.md5()
    }

  def Read(self, length=10000000):
    """Public read interface.

    Readers from this endpoint will receive a gzip file streamed from the infd.

    Args:
      length: How many bytes to read.

    Returns:
      A string.
    """

    # Read infd until we have length available in self.buff_fd.
    while self.zipper and self.buff_fd.len < length:
      data = self.infd.read(min(self.remaining_bytes, self.BUFFER_SIZE))
      if not data and self.zipper:
        # infd is finished.
        self.zipper.flush()
        self.zipper.close()
        self.zipper = None
        break

      l = len(data)
      self.total_read += l
      self.remaining_bytes -= l
      for h in self.hashers.values():
        h.update(data)
      self.zipper.write(data)

    return self.buff_fd.ReadFromFront(length)

  read = Read

  def __iter__(self):
    while 1:
      data = self.read(self.BUFFER_SIZE)
      if not data:
        break

      yield data

  def HashObject(self):
    return crypto.Hash(
        sha256=self.hashers["sha256"].digest(),
        sha1=self.hashers["sha1"].digest(),
        md5=self.hashers["md5"].digest())


class GunzipWrapper(object):
  """Wraps a file like object and decompresses it."""
  BUFFER_SIZE = 1024 * 1024

  def __init__(self, outfd):
    self.outfd = outfd
    # 16 + zlib.MAX_WBITS indicates that we want to decompress a gzip
    # compressed stream. See
    # http://stackoverflow.com/questions/1838699/
    # how-can-i-decompress-a-gzip-stream-with-zlib.
    self.decompressor = zlib.decompressobj(16 + zlib.MAX_WBITS)

  def Write(self, data):
    self.outfd.write(self.decompressor.decompress(data))

  def Flush(self):
    self.outfd.flush()

  def Close(self):
    self.outfd.close()

  write = Write
  flush = Flush
  close = Close
