#!/usr/bin/python
# -*- coding: utf-8 -*-
"""Report statistics about the artifact collection."""

from __future__ import print_function
import time

from artifacts import definitions
from artifacts import reader


class ArtifactStatistics(object):
  """Generate and print statistics about artifact files."""

  def _PrintDictAsTable(self, src_dict):
    key_list = src_dict.keys()
    key_list.sort()

    print('|', end='')
    for key in key_list:
      print(' {0} |'.format(key), end='')
    print('')

    print('|', end='')
    for key in key_list:
      print(' :---: |', end='')
    print('')

    print('|', end='')
    for key in key_list:
      print(' {0} |'.format(src_dict[key]), end='')
    print('\n')

  def PrintOSTable(self):
    print('**Artifacts by OS**\n')
    self._PrintDictAsTable(self.os_counts)

  def PrintLabelTable(self):
    print('**Artifacts by label**\n')
    self._PrintDictAsTable(self.label_counts)

  def PrintSourceTypeTable(self):
    print('**Artifacts by type**\n')
    self._PrintDictAsTable(self.source_type_counts)

  def PrintSummaryTable(self):
    print("""

As of {0} the repository contains:

| **File paths covered** | **{1}** |
| :------------------ | ------: |
| **Registry keys covered** | **{2}** |
| **Total artifacts** | **{3}** |
""".format(
        time.strftime('%Y-%m-%d'), self.path_count, self.reg_key_count,
        self.total_count))

  def BuildStats(self):
    artifact_reader = reader.YamlArtifactsReader()
    self.source_type_counts = {}
    self.label_counts = {}
    self.os_counts = {}
    self.total_count = 0
    self.path_count = 0
    self.reg_key_count = 0

    for artifact_definition in artifact_reader.ReadDirectory('definitions'):
      if hasattr(artifact_definition, 'labels'):
        for label in artifact_definition.labels:
          self.label_counts[label] = self.label_counts.get(label, 0) + 1

      for source in artifact_definition.sources:
        self.total_count += 1
        source_type = source.type_indicator
        self.source_type_counts[source_type] = self.source_type_counts.get(
            source_type, 0) + 1

        if source_type == definitions.TYPE_INDICATOR_WINDOWS_REGISTRY_KEY:
          self.reg_key_count += len(source.keys)
        if source_type == definitions.TYPE_INDICATOR_WINDOWS_REGISTRY_VALUE:
          self.reg_key_count += len(source.key_value_pairs)
        if (source_type == definitions.TYPE_INDICATOR_FILE or
            source_type == definitions.TYPE_INDICATOR_DIRECTORY):
          self.path_count += len(source.paths)

        os_list = source.supported_os
        for os_str in os_list:
          self.os_counts[os_str] = self.os_counts.get(os_str, 0) + 1

  def PrintStats(self):
    """Build stats and print in MarkDown format."""
    self.BuildStats()
    self.PrintSummaryTable()
    self.PrintSourceTypeTable()
    self.PrintOSTable()
    self.PrintLabelTable()


def main():
  statsbuilder = ArtifactStatistics()
  statsbuilder.PrintStats()


if __name__ == '__main__':
  main()
