# This file is part of Viper - https://github.com/viper-framework/viper
# See the file 'LICENSE' for copying permission.

import os
import tempfile
import string as printstring  # string is being used as a var - easier to replace here

try:
    from scandir import walk
except ImportError:
    from os import walk

from viper.common.abstracts import Module
from viper.core.database import Database
from viper.core.session import __sessions__
from viper.core.storage import get_sample_path
from viper.common.constants import VIPER_ROOT

try:
    import yara
    HAVE_YARA = True
except ImportError:
    HAVE_YARA = False


class YaraScan(Module):
    cmd = 'yara'
    description = 'Run Yara scan'
    authors = ['nex']

    def __init__(self):
        super(YaraScan, self).__init__()
        subparsers = self.parser.add_subparsers(dest='subname')
        parser_scan = subparsers.add_parser('scan', help='Scan files with Yara signatures')
        parser_scan.add_argument('-r', '--rule', help='Specify a ruleset file path (default will run data/yara/index.yara)')
        parser_scan.add_argument('-a', '--all', action='store_true', help='Scan all stored files (default if no session is open)')
        parser_scan.add_argument('-t', '--tag', action='store_true', help='Tag Files with Rule Name (default is not to)')
        parser_scan.add_argument('-v', '--verbose', action='store_true', help='Output a detailed overview of the matches and found offsets')

        parser_rules = subparsers.add_parser('rules', help='Operate on Yara rules')
        parser_rules.add_argument('-e', '--edit', help='Open an editor to edit the specified rule')

        for folder in ['/usr/share/viper/yara', os.path.join(VIPER_ROOT, 'data/yara')]:
            if os.path.exists(folder):
                self.rule_path = folder
                break

    def scan(self):

        def string_printable(line):
            line = str(line)
            new_line = ''
            for c in line:
                if c in printstring.printable:
                    new_line += c
                else:
                    new_line += '\\x' + c.encode('hex')
            return new_line

        # This means users can just drop or remove rule files without
        # having to worry about maintaining the index.
        # TODO: make paths absolute.
        # TODO: this regenerates the file at every run, perhaps we
        # could find a way to optimize this.
        def rule_index():
            tmp_path = os.path.join(tempfile.gettempdir(), 'index.yara')
            with open(tmp_path, 'w') as rules_index:
                for rule_file in os.listdir(self.rule_path):
                    # Skip if the extension is not right, could cause problems.
                    if not rule_file.endswith('.yar') and not rule_file.endswith('.yara'):
                        continue
                    # Skip if it's the index itself.
                    if rule_file == 'index.yara':
                        continue

                    # Add the rule to the index.
                    line = 'include "{0}"\n'.format(os.path.join(self.rule_path, rule_file))
                    rules_index.write(line)

            return tmp_path

        arg_rule = self.args.rule
        arg_scan_all = self.args.all
        arg_tag = self.args.tag
        arg_verbose = self.args.verbose

        # If no custom ruleset is specified, we use the default one.
        if not arg_rule:
            arg_rule = rule_index()

        # Check if the selected ruleset actually exists.
        if not os.path.exists(arg_rule):
            self.log('error', "No valid Yara ruleset at {0}".format(arg_rule))
            return

        # Compile all rules from given ruleset.
        rules = yara.compile(arg_rule)
        files = []

        # If there is a session open and the user didn't specifically
        # request to scan the full repository, we just add the currently
        # opened file's path.
        if __sessions__.is_set() and not arg_scan_all:
            files.append(__sessions__.current.file)
        # Otherwise we loop through all files in the repository and queue
        # them up for scan.
        else:
            self.log('info', "Scanning all stored files...")

            db = Database()
            samples = db.find(key='all')

            for sample in samples:
                files.append(sample)

        for entry in files:
            if entry.size == 0:
                continue

            self.log('info', "Scanning {0} ({1})".format(entry.name, entry.sha256))

            # Check if the entry has a path attribute. This happens when
            # there is a session open. We need to distinguish this just for
            # the cases where we're scanning an opened file which has not been
            # stored yet.
            if hasattr(entry, 'path'):
                entry_path = entry.path
            # This should be triggered only when scanning the full repository.
            else:
                entry_path = get_sample_path(entry.sha256)

            # Check if the file exists before running the yara scan.
            if not os.path.exists(entry_path):
                self.log('error', "The file does not exist at path {0}".format(entry_path))
                return

            rows = []
            tag_list = []
            found = False
            for match in rules.match(entry_path):
                found = True
                # Add a row for each string matched by the rule.
                if arg_verbose:
                    for string in match.strings:
                        rows.append([match.rule, string_printable(string[1]), string_printable(string[0]), string_printable(string[2])])
                else:
                    self.log('item', match.rule)
                # Add matching rules to our list of tags.
                # First it checks if there are tags specified in the metadata
                # of the Yara rule.
                match_tags = match.meta.get('tags')
                # If not, use the rule name.
                # TODO: as we add more and more yara rules, we might remove
                # this option and only tag the file with rules that had
                # tags specified in them.
                if not match_tags:
                    match_tags = match.rule

                # Add the tags to the list.
                tag_list.append([entry.sha256, match_tags])

            if arg_verbose and rows:
                header = [
                    'Rule',
                    'String',
                    'Offset',
                    'Content'
                ]
                self.log('table', dict(header=header, rows=rows))
            # If we selected to add tags do that now.
            if found and arg_tag:
                db = Database()
                for tag in tag_list:
                    db.add_tags(tag[0], tag[1])

                # If in a session reset the session to see tags.
                if __sessions__.is_set() and not arg_scan_all:
                    self.log('info', "Refreshing session to update attributes...")
                    __sessions__.new(__sessions__.current.file.path)

    def rules(self):
        arg_edit = self.args.edit

        # Retrieve the list of rules and populate a list.
        rules = []
        count = 1
        for folder, folders, files in walk(self.rule_path):
            for file_name in files:
                rules.append([count, os.path.join(folder, file_name)])
                count += 1

        # If the user wants to edit a specific rule, loop through all of them
        # identify which one to open, and launch the default editor.
        if arg_edit:
            for rule in rules:
                if int(arg_edit) == rule[0]:
                    os.system('"${EDITOR:-nano}" ' + rule[1])
                    break
        # Otherwise, just print the list.
        else:
            self.log('table', dict(header=['#', 'Path'], rows=rules))
            self.log('', "")
            self.log('', "You can edit these rules by specifying --edit and the #")

    def run(self):
        super(YaraScan, self).run()
        if self.args is None:
            return

        if not HAVE_YARA:
            self.log('error', "Missing dependency, install yara")
            return

        if self.args.subname == 'scan':
            self.scan()
        elif self.args.subname == 'rules':
            self.rules()
        else:
            self.log('error', 'At least one of the parameters is required')
            self.usage()
