#!/usr/bin/env python

## --------------------------------------------------------------------------------------------------------------------
#
# Copyright (c) 2021, rs-develop (rsdevelop.contact@gmail.com)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
## --------------------------------------------------------------------------------------------------------------------
#
# File:             ioccrawler.py
# Description:      The forensic ioc crawler extract iocs from directorys, files and mount points.
# Usage:            ioccrawler.py [-h]
# Author:           rs-develop
#
## --------------------------------------------------------------------------------------------------------------------

import logging         # for output
import os			   # for path, exit
import sys			   # for exit
import argparse		   # for program arguments
import timeit		   # for run time

from configparser import ConfigParser, ExtendedInterpolation# for loading config files
from crawler import crawler
from crawler.crawlererr import CralwerConfigAttributeError, CrawelrSetNewConfigError, CrawlerError, CrawlerConfigError, CrawlerPatternError
from crawler.exporter import CrawlerExporter

## --------------------------------------------------------------------------------------------------------------------

# init root logger
logging.root.setLevel(logging.NOTSET)

# Create logger
LOG = logging.getLogger('IocCrawlerLog')
LOG.setLevel(logging.NOTSET)

## --------------------------------------------------------------------------------------------------------------------
def convertTime(seconds):
    seconds = seconds % (24 * 3600)
    hour = seconds // 3600
    seconds %= 3600
    minutes = seconds // 60
    seconds %= 60
      
    return "%d:%02d:%02d" % (hour, minutes, seconds)

## --------------------------------------------------------------------------------------------------------------------
if __name__ == "__main__":

    try:

        # attributes
        printToStdout           = False
        basedir                 = os.path.dirname(os.path.abspath(__file__))
        working_directory       = os.path.join(basedir, os.path.dirname(crawler.__file__)) + "/"
        result_columns          = []
        pattern_columns         = []
        pattern_file            = ""
        whitelist_file          = ""
        user_pattern_loaded     = False
        user_whitelist_loaded   = False

        # read the crawler config
        config_file    = working_directory + 'data/config.ini'
        
        # read config
        try:
            config = ConfigParser(interpolation=ExtendedInterpolation())
            config.read(config_file)
            
            # read the result columns for printing to stdout
            result_columns.extend(config['settings']['result_columns'].strip().split('\n'))

            # Check if an individual user pattern or withelist file is set
            try:
                if config['settings']['default_user_pattern']:
                    if os.path.exists(config['settings']['default_user_pattern']):
                        pattern_file = config['settings']['default_user_pattern']
                        user_pattern_loaded = True
                    else:
                        raise CralwerConfigAttributeError('default_user_pattern', config_file, 'Switching to the default pattern file.')
                if config['settings']['default_user_whitelist']:
                    if os.path.exists(config['settings']['default_user_whitelist']):
                        whitelist_file = config['settings']['default_user_whitelist']
                        user_whitelist_loaded = True
                    else:
                        raise CralwerConfigAttributeError('default_user_whitelist', config_file, 'Switching to the default whitelist file.')

            except CralwerConfigAttributeError as cae:
                print(cae.msg)
            
            # try to load default pattern and whitelist file
            if not user_pattern_loaded:
                if os.path.exists(working_directory + config['settings']['default_pattern']):
                    pattern_file = working_directory + config['settings']['default_pattern']
                else:
                    raise CrawlerConfigError('Error while loading the default pattern file from: ' + config['settings']['default_pattern'] + '. You can download it agian from: https://github.com/rs-develop/ForIocCrawler/blob/main/crawler/data/pattern.ini')

            # try to load default whitelist file
            if not user_whitelist_loaded:
                if os.path.exists(working_directory + config['settings']['default_whitelist']):
                    whitelist_file = working_directory + config['settings']['default_whitelist']
                else:
                    raise CrawlerConfigError('Error while loading the default pattern file from: ' + config['settings']['default_whitelist'] + '. You can download it agian from: https://github.com/rs-develop/ForIocCrawler/blob/main/crawler/data/whitelist.ini')

        except Exception as e:
            raise CrawlerConfigError(getattr(e, 'message', repr(e)))

        # read sections from pattern
        try:
            pattern = ConfigParser(interpolation=ExtendedInterpolation())
            pattern.read(pattern_file)
            pattern_columns = [x.lower() for x in pattern.sections()]
        except Exception as e:
            raise CrawlerConfigError(getattr(e, 'message', repr(e)))

        # parse arguments
        parser = argparse.ArgumentParser(description="IoC crawler for parsing files, directories or mount points.", usage='%(prog)s -h/--help for help')

        # Create Subparser
        subparsers = parser.add_subparsers(help='sub-command help')

        # Create Subparser for the ioc crawler
        ioc_crawler_parser = subparsers.add_parser('parse', help='Subcommand for ioc crawler')
        ioc_crawler_parser.add_argument('source_file_or_dir', help='Source file, directory or mount point.')
        ioc_crawler_parser.add_argument('-c', '--columns', choices=result_columns + ["all"], default="all", nargs='+', type=str.lower, help='Output columns. All columns will be printed by default.')
        ioc_crawler_parser.add_argument('-t', '--type', nargs="+", default="all", help="Print only the provided pattern types. Available pattern types are: \"%s\"" %" ".join([x.lower() for x in pattern_columns]))
        ioc_crawler_parser.add_argument('-m', '--mode', choices=["stdout","forensics"], default='stdout', help='Output mode. Print results to stdout (default) or run in forensics mode with processing status and summary.')
        ioc_crawler_parser.add_argument('-o', '--out',  dest='output_file_name', help='Output file name (works also in stdout mode).')
        ioc_crawler_parser.add_argument('-w', '--whitelist', action='store_true', help="Enables whitelisting. Use \".forioccrawler config --print-whitelist\" to view the content of the whitelist.")
        ioc_crawler_parser.add_argument('--load-whitelist', dest='individual_whitelist_file', help='Use a individual whitelist. For changing the whitelist file permanently, see the config subcommand.')
        ioc_crawler_parser.add_argument('--load-pattern', dest='individual_pattern_file', help='Use a individual pattern file. For changing the pattern file permanently, see the config subcommand.')
        ioc_crawler_parser.add_argument('--threads', type=int, default=int(config['settings']['default_process_count']), help='Set the thread count for the parsing. To change the default, check the config subcommand. (default=%d, max=%d)' % (int(config['settings']['default_process_count']), int(config['settings']['max_processes'])))
        ioc_crawler_parser.add_argument('-n', action='store_false', dest='match_highlighting', default=True, help="No match highligting")
        ioc_crawler_parser.add_argument('-s', dest='match_size', default=256, type=int, help="Set maximal match size (default=256). Have to be greater then 5.")
        ioc_crawler_parser.add_argument("-v", "--verbose", action = "store_true", help='Show debug messages and write debug log')
        ioc_crawler_parser.add_argument("--time", action = "store_true", help='Show run time.')

        # Create Subparser for config
        config_parser = subparsers.add_parser('config', help='Subcommand for configuration informations')
        config_parser.add_argument('--show', action='store_true', help='Shows the content of the configuration file')
        config_parser.add_argument('--set-pattern', dest='user_pattern_file', help='Set a personal/individual pattern file as new default.')
        config_parser.add_argument('--set-whitelist', dest='user_whitelist_file', help='Set a personal/individual whitelist file as new default.')
        config_parser.add_argument('--print-whitelist', action='store_true', help='Prints the path and the content of the default whitelist and exits.')
        config_parser.add_argument('--print-pattern', action='store_true', help='Prints the path and the content of the default pattern file and exits.')
        config_parser.add_argument('--set-thread-count', dest='new_thread_count', help='Change the default thread count for processing. (default=%d, max=%d)' % (int(config['settings']['default_process_count']), int(config['settings']['max_processes'])))
        config_parser.add_argument('--restore-pattern', action='store_true', help='Restores the default pattern file for parsing.')
        config_parser.add_argument('--restore-whitelist', action='store_true', help='Restores the default whitelist file for whitelisting.')
        config_parser.add_argument('--reset-config', action='store_true', help='Resets all config settings to default.')
        
        # Create Subparser for version
        version_parser = subparsers.add_parser('version', help='Subcommand for version information')
        version_parser.add_argument('--show', action='store_true', help='Show program version')
        
        # Parse Arguments
        args = parser.parse_args()

        # If no argument is provided, show help and exit
        if len(sys.argv) < 2:
            parser.print_help()
            exit()

        ## -------------------------------------------------------------------
        ## Subcommand config
        # If the config argument is provided, show configs
        if 'config' in sys.argv:
            command_passed = False
            
            # Check if any argument is passed
            for key in args.__dict__:
                if args.__dict__[key] is not None and args.__dict__[key] is not False:
                    command_passed = True
                    break
            if command_passed:
                if args.show:
                    print('default process count' + ' : ' + config['settings']['default_process_count'])
                    print('max process count' + ' : ' + config['settings']['max_processes'])
                    print('default pattern file' + ' : ' + os.path.abspath(config['settings']['default_pattern']))
                    print('default whitelist file' + ' : ' + os.path.abspath(config['settings']['default_whitelist']))
                    print('individual user pattern' + ' : ' + config['settings']['default_user_pattern'])
                    print('individual user whitelist' + ' : ' + config['settings']['default_user_whitelist'])
                if args.user_pattern_file:
                    if (os.path.exists(os.path.abspath(args.user_pattern_file))):
                        try:
                            config['settings']['default_user_pattern'] = os.path.abspath(args.user_pattern_file)
                            # try to read the new config
                            test_config = ConfigParser(interpolation=ExtendedInterpolation())
                            test_config.read(config['settings']['default_user_pattern'])
                            with open(config_file, 'w') as updated_configuration:
                                config.write(updated_configuration)
                            print('[+] Updated configuration.')
                        except Exception as e:
                            raise CrawelrSetNewConfigError(args.user_pattern_file, getattr(e, 'message', repr(e)))
                        exit()
                if args.user_whitelist_file:
                    if (os.path.exists(os.path.abspath(args.user_whitelist_file))):
                        try:
                            config['settings']['default_user_whitelist'] = os.path.abspath(args.user_whitelist_file)
                            # try to read the new config
                            test_config = ConfigParser(interpolation=ExtendedInterpolation())
                            test_config.read(config['settings']['default_user_whitelist'])
                            with open(config_file, 'w') as updated_configuration:
                                config.write(updated_configuration)
                            print('[+] Updated configuration.')
                        except Exception as e:
                            raise CrawelrSetNewConfigError(args.user_whitelist_file, getattr(e, 'message', repr(e)))
                        exit()
                if args.print_pattern:
                    print("Path:", pattern_file)
                    with open(pattern_file,'r') as f:
                        print(f.read())
                    exit()
                if args.print_whitelist:
                    print("Path:", pattern_file)
                    with open(whitelist_file,'r') as f:
                        print(f.read())
                    exit()
                if args.new_thread_count:
                    if int(args.new_thread_count) > 1 and int(args.new_thread_count) < int(config['settings']['max_processes']):
                        config['settings']['default_process_count'] = args.new_thread_count
                        with open(config_file, 'w') as updated_configuration:
                                config.write(updated_configuration)
                        print('[+] Updated configuration.')
                        exit()
                if args.restore_pattern:
                    config['settings']['default_user_pattern'] = ""
                    with open(config_file, 'w') as updated_configuration:
                        config.write(updated_configuration)
                    print('[+] Updated configuration.')
                    exit()
                if args.restore_whitelist:
                    config['settings']['default_user_whitelist'] = ""
                    with open(config_file, 'w') as updated_configuration:
                        config.write(updated_configuration)
                    print('[+] Updated configuration.')
                    exit()
                if args.reset_config:
                    config['settings']['default_process_count'] = '4'
                    config['settings']['default_user_whitelist'] = ""
                    config['settings']['default_user_pattern'] = ""
                    with open(config_file, 'w') as updated_configuration:
                        config.write(updated_configuration)
                    print('[+] Updated configuration.')
                    exit()
            # if no argument is passed, show help
            else:
                config_parser.print_help()
        ## -------------------------------------------------------------------

        ## -------------------------------------------------------------------
        ## Subcommand parser
        ## if parse is provided, check the arguments and execute the parser
        elif 'parse' in sys.argv:

            # if verbose flag is set, set logging to debug
            if args.verbose:
                LOG.setLevel(logging.DEBUG)

                # Create debug log handler for console
                debug_log = logging.StreamHandler()
                debug_log.setLevel(logging.DEBUG) # log only debug
                debug_log.setFormatter(logging.Formatter('DEBUG %(asctime)-15s %(processName)s %(module)s %(lineno)d %(message)s'))

                # Create file log handler for debug
                if os.path.exists('debug.log'):
                    os.remove('debug.log')
                fh_debug = logging.FileHandler('debug.log')
                fh_debug.setLevel(logging.DEBUG) # log only debug
                fh_debug.setFormatter(logging.Formatter('DEBUG %(asctime)-15s %(processName)s %(module)s %(lineno)d %(message)s'))

                # Add log handler
                LOG.addHandler(debug_log)
                LOG.addHandler(fh_debug)
                
                # set timer for run time
                start = timeit.default_timer()

                LOG.debug("Debug initialisation finished")
                LOG.debug("Arguments: %s" %str(args))

            # if time arg, show run time at the end
            elif args.time:
                start = timeit.default_timer()

            # check program mode
            if args.mode not in ["forensics", "stdout"]:
                ioc_crawler_parser.print_help()
                raise CrawlerError("Unknown program mode: %s." %args.mode)
            elif args.mode == "forensics" and args.output_file_name == None:
                ioc_crawler_parser.print_help()
                raise CrawlerError("Forensics-Mode requires an output file. (try -o)")

            # check for stdout option and reset output path
            if args.mode == "stdout":
                printToStdout = True

            # check format options for reslult columns
            if args.columns != "all":
                result_columns = args.columns

            # if an individual pattern file is set, check pattern
            if args.individual_pattern_file:
                try:
                    pattern = ConfigParser(interpolation=ExtendedInterpolation())
                    pattern.read(args.individual_pattern_file)
                    pattern_columns = [x.lower() for x in pattern.sections()]
                except Exception as e:
                    raise CrawlerPatternError(getattr(e, 'message', repr(e)))
                
                pattern_file = args.individual_pattern_file

            # check if selected pattern type is "all", then set all pattern from the individual pattern file
            if args.type == "all":
                args.type = pattern_columns
            # if not all is set, check if the pattern is present in the pattern file, if not: throw exception
            else:
                for pattern in args.type:
                        if pattern not in pattern_columns:
                            raise CrawlerPatternError("Unknown pattern %s in %s" %(pattern, pattern_file))

            # if whitelisting is enabled set the path to the default or individual whitelist
            if args.whitelist or args.individual_whitelist_file:

                # if an individual whitelist file is set, change path
                if args.individual_whitelist_file:
                    whitelist_file = args.individual_whitelist_file
            # If whitelisting is disabled, clear the path
            else:
                whitelist_file = ""

            ## -------------------------------------------------------------------
            # run crawler
            ioccrawler = crawler.Crawler(args.source_file_or_dir, args.threads, pattern_file, printToStdout, result_columns, 
                                        args.type, args.match_highlighting, args.match_size, whitelist_file, 0, 0)
            ioccrawler.do()
            ## -------------------------------------------------------------------

            # check the export option
            if args.output_file_name:
                print('[+] Writing Export')
                CrawlerExporter.csvExport(args.output_file_name, ioccrawler.resultList, result_columns)
                print('[+] Results written to: %s' %(args.output_file_name))

            # show run time
            if args.time or args.verbose:
                stop = timeit.default_timer()
                print('[+] Jobs finished in: %s (H:MM:SS)' %(convertTime(stop - start)))

            # show summary if not printed to stdout
            if not printToStdout:
                
                summary = ioccrawler.getResultSummary()
                print("[+] Summary of matches")
                for ioc in summary:
                    print(" |- %s: %s" %(ioc, summary[ioc]))
            
                # show run time
                if args.time or args.verbose:
                    stop = timeit.default_timer()
                    print('[+] Jobs finished in: %s (H:MM:SS)' %(convertTime(stop - start)))
            
                print("[+] Done")
            # end if
        ## -------------------------------------------------------------------
        ## Subcommand version
        # Show Program Version
        elif 'show' in vars(args):
            print(config['settings']['version'])
            exit()
        ## -------------------------------------------------------------------

    # Handle errors
    except CrawlerError as err:
        print(err.msg)
    except KeyboardInterrupt:
        print("[!] User interrupt.")
        try:
            sys.exit(0)
        except SystemExit:
            os._exit(0)
    except Exception as e:
        print("[!] Unhandled error: " + str(e))

# end main