# -*- coding: utf-8 -*-
# Copyright 2015 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Task to analyse Hadoop AppRoot files."""

from __future__ import unicode_literals

import codecs
import logging
import os
import subprocess

from turbinia import TurbiniaException

from turbinia.lib import text_formatter as fmt
from turbinia.evidence import ReportText
from turbinia.lib.utils import extract_artifacts
from turbinia.workers import TurbiniaTask
from turbinia.workers import Priority

log = logging.getLogger('turbinia')


class HadoopAnalysisTask(TurbiniaTask):
  """Task to analyse Hadoop AppRoot files."""

  def _AnalyzeHadoopAppRoot(self, collected_artifacts, output_dir):
    """Runs a naive AppRoot files parsing method.

    This extracts strings from the saved task file, and searches for usual
    post-compromise suspicious patterns.

    TODO: properly parse the Proto. Some documentation can be found over there:
    https://svn.apache.org/repos/asf/hadoop/common/branches/branch-0.23.7/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/proto/yarn_protos.proto

    Args:
      collected_artifacts(list(str)): a list of paths to extracted files
      output_dir(str): The base directory the artfacts are in.

    Returns:
      Tuple(
        list(str): The report data as a list of lines
        report_priority(int): The priority of the report (0 - 100)
        summary(str): A summary of the report (used for task status)
      )
    """
    report = []
    evil_commands = []
    strings_count = 0
    priority = Priority.MEDIUM
    summary = ''
    for filepath in collected_artifacts:
      relpath = os.path.relpath(filepath, output_dir)
      command = 'strings -a "{0:s}"'.format(filepath)
      log.debug('Running command [{0:s}]'.format(command))
      proc = subprocess.Popen(
          command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
      strings_output, _ = proc.communicate()
      strings_output = codecs.decode(strings_output, 'utf-8')
      for line in strings_output.splitlines():
        strings_count += 1
        if (line.find('curl') >= 0) or (line.find('wget') >= 0):
          evil_commands.append((relpath, line))

    if evil_commands:
      msg = 'Found suspicious commands!'
      report.append(fmt.heading4(fmt.bold(msg)))
      summary = msg
      priority = Priority.CRITICAL
    else:
      msg = 'Did not find any suspicious commands.'
      report.append(fmt.heading4(msg))
      summary = msg

    for filepath, command in evil_commands:
      report.append(fmt.bullet(fmt.bold('Command:')))
      report.append(fmt.code(command))
      report.append('Found in file:')
      report.append(fmt.code(filepath))

    msg = 'Extracted {0:d} strings from {1:d} file(s)'.format(
        strings_count, len(collected_artifacts))
    report.append(fmt.bullet(msg))

    return (report, priority, summary)

  def run(self, evidence, result):
    """Run Hadoop specific analysis on the evidences.

    Args:
        evidence (Evidence object):  The evidence we will process
        result (TurbiniaTaskResult): The object to place task results into.

    Returns:
        TurbiniaTaskResult object.
    """

    # What type of evidence we should output.
    output_evidence = ReportText()

    # Where to store the resulting output file.
    output_file_name = 'hadoop_analysis.txt'
    output_file_path = os.path.join(self.output_dir, output_file_name)

    output_evidence.local_path = output_file_path

    try:
      # We don't use FileArtifactExtractionTask as it export one evidence per
      # file extracted
      output_dir = os.path.join(self.output_dir, 'artifacts')
      collected_artifacts = extract_artifacts(
          artifact_names=['HadoopAppRoot'], disk_path=evidence.local_path,
          output_dir=output_dir)

      (report, priority, summary) = self._AnalyzeHadoopAppRoot(
          collected_artifacts, output_dir)
      if not report:
        raise TurbiniaException(
            'Report generated by _AnalyzeHadoopAppRoot() is empty')

      output_evidence.text_data = '\n'.join(report)
      result.report_data = output_evidence.text_data

      # Write the report to the output file.
      with open(output_file_path, 'w') as fh:
        fh.write(output_evidence.text_data.encode('utf8'))
        fh.write('\n'.encode('utf8'))

      result.add_evidence(output_evidence, evidence.config)
      result.report_priority = priority
      result.close(self, success=True, status=summary)
    except TurbiniaException as e:
      result.close(self, success=False, status=str(e))
      return result
    return result
