#!/usr/bin/env python
# Lint as: python3
"""Plugin that exports results as SQLite db scripts."""

import collections
import io
import os
import zipfile

import sqlite3

from grr_response_core.lib import rdfvalue
from grr_response_core.lib import utils
from grr_response_core.lib.rdfvalues import structs as rdf_structs
from grr_response_core.lib.util import collection
from grr_response_core.lib.util.compat import yaml
from grr_response_server import instant_output_plugin


class Rdf2SqliteAdapter(object):
  """An adapter for converting RDF values to a SQLite-friendly form."""

  class Converter(object):

    def __init__(self, sqlite_type, convert_fn):
      self.sqlite_type = sqlite_type
      self.convert_fn = convert_fn

  BYTES_CONVERTER = Converter("BLOB", bytes)
  STR_CONVERTER = Converter("TEXT", str)

  DEFAULT_CONVERTER = STR_CONVERTER

  INT_CONVERTER = Converter("INTEGER", int)

  # Converters for fields that have a semantic type annotation in their
  # protobuf definition.
  SEMANTIC_CONVERTERS = {
      rdfvalue.RDFString:
          STR_CONVERTER,
      rdfvalue.RDFBytes:
          BYTES_CONVERTER,
      rdfvalue.RDFInteger:
          INT_CONVERTER,
      bool:
          INT_CONVERTER,  # Sqlite does not have a bool type.
      rdfvalue.RDFDatetime:
          Converter("INTEGER", lambda x: x.AsMicrosecondsSinceEpoch()),
      rdfvalue.RDFDatetimeSeconds:
          Converter("INTEGER", lambda x: x.AsSecondsSinceEpoch() * 1000000),
      rdfvalue.DurationSeconds:
          Converter("INTEGER", lambda x: x.microseconds),
  }

  # Converters for fields that do not have a semantic type annotation in their
  # protobuf definition.
  NON_SEMANTIC_CONVERTERS = {
      rdf_structs.ProtoBinary: BYTES_CONVERTER,
      rdf_structs.ProtoString: STR_CONVERTER,
      rdf_structs.ProtoEnum: STR_CONVERTER,
      rdf_structs.ProtoUnsignedInteger: INT_CONVERTER,
      rdf_structs.ProtoSignedInteger: INT_CONVERTER,
      rdf_structs.ProtoFixed32: INT_CONVERTER,
      rdf_structs.ProtoFixed64: INT_CONVERTER,
      rdf_structs.ProtoFloat: Converter("REAL", float),
      rdf_structs.ProtoDouble: Converter("REAL", float),
      rdf_structs.ProtoBoolean: INT_CONVERTER,
  }

  @staticmethod
  def GetConverter(type_info):
    if type_info.__class__ is rdf_structs.ProtoRDFValue:
      return Rdf2SqliteAdapter.SEMANTIC_CONVERTERS.get(
          type_info.type, Rdf2SqliteAdapter.DEFAULT_CONVERTER)
    else:
      return Rdf2SqliteAdapter.NON_SEMANTIC_CONVERTERS.get(
          type_info.__class__, Rdf2SqliteAdapter.DEFAULT_CONVERTER)


class SqliteInstantOutputPlugin(
    instant_output_plugin.InstantOutputPluginWithExportConversion):
  """Instant output plugin that converts results into SQLite db commands."""

  plugin_name = "sqlite-zip"
  friendly_name = "SQLite scripts (zipped)"
  description = "Output ZIP archive containing SQLite scripts."
  output_file_extension = ".zip"

  ROW_BATCH = 100

  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.archive_generator = None  # Created in Start()
    self.export_counts = {}

  @property
  def path_prefix(self):
    prefix, _ = os.path.splitext(self.output_file_name)
    return prefix

  def Start(self):
    self.archive_generator = utils.StreamingZipGenerator(
        compression=zipfile.ZIP_DEFLATED)
    self.export_counts = {}
    return []

  def ProcessSingleTypeExportedValues(self, original_value_type,
                                      exported_values):
    first_value = next(exported_values, None)
    if not first_value:
      return

    if not isinstance(first_value, rdf_structs.RDFProtoStruct):
      raise ValueError("The SQLite plugin only supports export-protos")
    yield self.archive_generator.WriteFileHeader(
        "%s/%s_from_%s.sql" % (self.path_prefix, first_value.__class__.__name__,
                               original_value_type.__name__))
    table_name = "%s.from_%s" % (first_value.__class__.__name__,
                                 original_value_type.__name__)
    schema = self._GetSqliteSchema(first_value.__class__)

    # We will buffer the sql statements into an in-memory sql database before
    # dumping them to the zip archive. We rely on the PySQLite library for
    # string escaping.
    db_connection = sqlite3.connect(":memory:")
    db_cursor = db_connection.cursor()

    yield self.archive_generator.WriteFileChunk(
        "BEGIN TRANSACTION;\n".encode("utf-8"))

    with db_connection:
      buf = io.StringIO()
      buf.write(u"CREATE TABLE \"%s\" (\n  " % table_name)
      column_types = [(k, v.sqlite_type) for k, v in schema.items()]
      buf.write(u",\n  ".join([u"\"%s\" %s" % (k, v) for k, v in column_types]))
      buf.write(u"\n);")
      db_cursor.execute(buf.getvalue())

      chunk = (buf.getvalue() + "\n").encode("utf-8")
      yield self.archive_generator.WriteFileChunk(chunk)

      self._InsertValueIntoDb(table_name, schema, first_value, db_cursor)

    for sql in self._FlushAllRows(db_connection, table_name):
      yield sql
    counter = 1
    for batch in collection.Batch(exported_values, self.ROW_BATCH):
      counter += len(batch)
      with db_connection:
        for value in batch: