import subprocess
import os
import csv
import sys
from sys import platform as _platform
import traceback
import argparse
import re

# Platform
os_platform = ""

if _platform == "linux" or _platform == "linux2":
    os_platform = "linux"
elif _platform == "darwin":
    os_platform = "macos"
elif _platform == "win32":
    os_platform = "windows"

# Win32 Imports
if os_platform == "windows":
    try:
        import wmi
        import win32api
        from win32com.shell import shell
        import win32file
    except Exception as e:
        print("Linux System - deactivating process memory check ...")
        os_platform = "linux"  # crazy guess

if os_platform == "":
    print("Unable to determine platform - LOKI is lost.")
    sys.exit(1)

class Node:
    def __init__(self, pid, ppid, image_file_name):
        self.pid = pid
        self.ppid = ppid
        self.image_file_name = image_file_name
        self.children = []

class AutoVol3(object):

    normal_proc = {}
    suspect_proc = {}

    regex_patterns = []
    whitelist_paths = []
    normal_paths =[]
    normal_paths_x86 =[]        # The normal path for a process if x86
    normal_process_path = []    # The normal path for a process
    process_path = {}           # The path for the running process from dlllist
    normal_sids = {}

    suspecious_proc_sids = {}
    suspect_cmdlines = {}
    suspect_netscan = {}

    dll_stacking = {}

    process_tree = {}
    pslist = {}
    psscan ={}
    dlllist = {}
    cmdline = {}
    netscan ={}
    getsids ={}
    handles = {}

    baseline_proc = {}
    evidence_bag = []

    score = {}

    plugins = ['pslist', 'psscan', 'pstree', 'dlllist', 'cmdline', 'netstat', 'netscan', 'handle', 'getsids']

    analyze = ['List suspecious processes from EPROCESS', 'List suspecious processes with memory scanning', 'List suspecious cmdlines', 'List suspecious DLLs',
               'List suspecious communications from EPROCESS', 'List suspecious communications with memory scanning', 'Compare the suspect image with baseline',
               'List processes running with suspecious account']

    csv_files = ['pslist.csv', 'psscan.csv', 'pstree.csv', 'dlllist.csv', 'cmdline.csv', 'netstat.csv', 'netscan.csv', 'handles.csv', 'getsids.csv']

    def __init__(self, image):

        self.image = image
        self.app_path = get_application_path()

        self.csvgen(os.path.join(self.app_path, args.p))
        self.initialize_normal_proc(os.path.join(self.app_path, 'normal_proc.txt'))
        self.initialize_normal_paths(os.path.join(self.app_path, 'normal_paths.txt'))
        self.initialize_normal_sids(os.path.join(self.app_path, 'normal_sids.txt'))
        self.initialize_whitelist_paths('whitelist.txt')
        self.initialize_regex_patterns('regex_patterns.txt')

        # Read the required file from running the volatility plugins.
        self.csv_reader('pslist.csv')
        self.csv_reader('psscan.csv')
        self.csv_reader('dlllist.csv')
        self.csv_reader('cmdline.csv')
        self.csv_reader('netscan.csv')
        self.csv_reader('getsids.csv')
        self.csv_reader('handles.csv')
        self.malhandles()

        # Analyzing memory for suspicious processes
        self.procTree(os.path.join(self.app_path, args.p) + '/'+'pslist.csv')
        self.malproc(self.process_tree['4'])
        self.malpath()
        self.malcmdline()
        self.malcomm('blacklist.txt')
        self.baseline('proc_baseline.txt')
        self.malgetsids()
        self.find_processes_without_parents()

        # The following function must be run the last one
        self.malicious_weight()

        header = ['Source Name', 'PID', 'PPID','Process Name', 'Path', 'Timestamps', 'Long Description']
        self.evidence_bag.append(','.join(header))

    def run_plugin(self, memory_image_path, output_file):

        # if output_file == 'netscan.csv':
        #     command = ["vol2.py", "--profile", "Win10x64_15063", "-f", memory_image_path, output_file[:-4]]

        # else:
        command = ["vol", "-r", "csv", "-f", memory_image_path, 'windows.'+output_file[:-4]]

        try:
            result = subprocess.run(command, capture_output=True, text=True, check=True)
            output = result.stdout
            # Save the result to the specified CSV file
            with open(os.path.join(self.app_path, args.p)+'/'+output_file, 'w', encoding='utf-8') as csv_file:
                csv_file.write(output)

        except subprocess.CalledProcessError as e:
            print("Error:", e)
            print("Return Code:", e.returncode)
            print("Command Output:", e.output)
        except Exception as e:
            print("An error occurred:", str(e))

    def csvgen(self, memory_image_path):

        # Specify the directory path you want to list files from
        #directory_path = "./memory"

        try:

            # Use the os.listdir() function to get a list of files in the directory
            files = os.listdir(memory_image_path)

            # Check if there are any .img files in the directory
            img_files = [file for file in files if file.endswith(".img")]

            if not img_files:
                print(f"No '.img' files found in '{memory_image_path}'.")
            else:
                # Loop through and print the names of .img fil

                for file in self.csv_files:

                    if file not in files:
                        print('File not found:', file)
                        print('Generate file:', file,'...')
                        self.run_plugin(memory_image_path+'/'+img_files[0],file)

                        # if file == 'dlllist.csv':
                        #     result = list_dlls(img_files[0])
                        # elif file == 'cmdline.csv':

        except FileNotFoundError:
            print(f"The directory '{memory_image_path}' does not exist.")
        except PermissionError:
            print(f"You do not have permission to access '{memory_image_path}'.")
        except Exception as e:
            print(f"An error occurred: {str(e)}")

    def baseline(self, baseline_file):
        try:
            # Use the os.listdir() function to get a list of files in the directory
            fpath = os.path.join(self.app_path, args.p)
            files = os.listdir(fpath)

            # If proc_baseline.txt is not found, uncomment the next two lines to create it. It will takkke sometime. Be patient
            if baseline_file not in files:
                # Define the command as a list of strings
                print('proc_baseline.txt file was found, generating this file ...')
                command = ["python3", "/opt/memory-baseliner/baseline.py", "-proc", "-i", "/cases/memory/base-rd01-memory.img",
                "--loadbaseline", '--jsonbaseline', '/cases/precooked/memory/Win10x64_proc.json', '-o', fpath+'/'+'proc_baseline.txt']
                # Run the command and capture its output
                completed_process = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

        except FileNotFoundError:
            print(f"The directory '{fpath}' does not exist.")
        except PermissionError:
            print(f"You do not have permission to access '{fpath}'.")
        except Exception as e:
            print(f"An error occurred: {str(e)}")

        # Read the content of the input file and replace '|' with ','
        with open(fpath+'/'+'proc_baseline.txt', "r") as input_file:
            content = input_file.read()
            modified_content = content.replace("|", ",")

            modified_content = modified_content.strip().split('\n')[1:]

            for row in modified_content:
                columns = row.split(',')
                pid = columns[0]

                if '.exe' in columns[6]:
                    # if pid not in self.baseline_proc:
                    #     self.baseline_proc[pid] = []
                    self.baseline_proc[pid.strip('"')] = [''.join(row)]
        #print(self.baseline_proc)
        # Write the modified content back to the output file
        # with open(fpath+'/'+'suspecious_proc.txt', "w") as output_file:
        #     output_file.write(modified_content)

        # print("Replacement complete. Output written to", 'suspecious_proc.txt')

    def initialize_normal_proc(self, normal_proc_file):

        # Open and read the "normal_proc.txt" file
        with open(normal_proc_file, 'r') as file:
            for line in file:
                # Split each line into parent and child using ':' as the separator
                parent, child = line.strip().split(':')

                # If the parent is not in the dictionary, add it with an empty list of children
                if parent not in self.normal_proc:
                    self.normal_proc[parent] = []

                # Append the child to the parent's list of children
                self.normal_proc[parent].append(child)

        # The resulting dictionary is parsed_normal_proc

    def initialize_whitelist_paths(self, whitelist_file):

        # Open and read the "normal_proc.txt" file
        # You need to change this list "normal_proc.txt", because it is not clean, missing a lot of system files and dlls.
        with open(whitelist_file, 'r') as file:
            for line in file:
                # Split each line into parent and child using ':' as the separator
               self.whitelist_paths.append(line.strip('\n').split(',')[0].lower())
               #print(self.whitelist_paths)

    def initialize_normal_sids(self, normal_sids_file):

        # Open and read the "normal_proc.txt" file
        with open(normal_sids_file, 'r') as file:
            for line in file:
                # Split each line into parent and child using ':' as the separator
                parent, child = line.strip().split(':')

                # If the parent is not in the dictionary, add it with an empty list of children
                if parent not in self.normal_sids:
                    self.normal_sids[parent] = []

                # Append the child to the parent's list of children
                self.normal_sids[parent].append(child)

    def initialize_normal_paths(self, fpath):
        try:

            with open(fpath, newline='') as file:

                for line in file:
                   # self.normal_process_path.append(line.strip('\n').split('\\')[-1])

                    if 'program files (x86)' in line:
                        self.normal_paths_x86.append(line.strip('\n'))
                    else:
                        self.normal_paths.append(line.strip('\n'))

        except FileNotFoundError:
            print(f"The directory '{fpath}' does not exist.")
        except PermissionError:
            print(f"You do not have permission to access '{fpath}'.")
        except Exception as e:
            print(f"An error occurred: {str(e)}")

    def initialize_blacklist_addresses(self, file_path):
    # Read the blacklist addresses from a file and return them as a set
        with open(file_path, mode='r') as blacklist_file:
            return set(line.strip() for line in blacklist_file)

    def initialize_regex_patterns(self, regex_file):
        # Read regex patterns from the file

        with open(regex_file, 'r') as patterns_file:
            for line in patterns_file:
                if line.startswith('#') or not line:
                    continue                        # Skip this line and move to the next one

                line = line.strip()  # Remove leading/trailing whitespace
                if line:
                    self.regex_patterns.append(line)

        #print(self.regex_patterns)

    def procTree(self, proc_file):
        try:

            with open(proc_file, newline='') as csvfile:
                reader = csv.DictReader(csvfile)

                for row in reader:
                    pid = row['PID']
                    ppid = row['PPID']
                    image_file_name = row['ImageFileName']

                    # Create a new node if it doesn't exist
                    if pid not in self.process_tree:
                        self.process_tree[pid] = Node(pid, ppid, image_file_name)

                    # Create a new node for the parent (PPID) if it doesn't exist
                    if ppid not in self.process_tree:
                        self.process_tree[ppid] = Node(ppid, 0, '')

                    # Update the image file name for the current node
                    self.process_tree[pid].image_file_name = image_file_name

                    # Add the current node as a child to its parent
                    add_child(self.process_tree[ppid], self.process_tree[pid])

        except FileNotFoundError:
            print(f"The directory '{proc_file}' does not exist.")
        except PermissionError:
            print(f"You do not have permission to access '{proc_file}'.")
        except Exception as e:
            print(f"An error occurred: {str(e)}")

    def get_parent(self, pid):
        if pid in self.process_tree:
            node = self.process_tree[pid]
            parent_pid = node.ppid
            if parent_pid in self.process_tree:
                return self.process_tree[parent_pid]
        return None

    def find_parent_recursive(self, child_pid):
        # Get the parent node for the given child_pid
        parent_node = self.get_parent(child_pid)

        # Check if a parent node was found
        if parent_node:

            if parent_node.pid not in self.score:
                #print(f"Parent PID: {parent_node.pid}, Image File Name: {parent_node.image_file_name}")

                # Because score keeps track of pid while evidence_bag not. In addition, let the score of the parent is equal to child
                self.score[parent_node.pid] = self.score[child_pid]
                self.collect_evidence(parent_node.pid )
            # If the parent is not services.exe, recursively find its parent
            if parent_node.image_file_name != 'svchost.exe':
                self.find_parent_recursive(parent_node.pid)
        # else:
        #     print(f"No parent found for PID {child_pid}.")

    def malproc(self,node):

        # 1. Should find abnormal parent-child relationship                                                             ==> Done
        # 2. Should find zero parent processes                                                                          ==> Done
        # 3. Should find processes running from weired paths                                                            ==> Done
        # 4. Should hidden/Unlinked processes like processes found in psscan, but not in pslist                         ==> Done
        # 5. When does the process start?                                                                          ==>TBD
        # 6. Should find any process that impersonate known processes and trying to blind in with normal processes      ==> Done

        # 1. Should find abnormal parent-child relationship

        if not node:
            return #suspect_proc

        for child in node.children:
            if not(node.image_file_name in self.normal_proc and child.image_file_name in self.normal_proc[node.image_file_name]):
                # print(node.image_file_name, child.image_file_name)
                self.suspect_proc[child.pid] =self.pslist[child.pid][2:] + ', Suspecious parent-child relationship'

            self.malproc(child)

        return #suspect_proc

    # 2. Should find zero parent processes

    def find_processes_without_parents(self):
        for pid, node in self.process_tree.items():

            if node.ppid not in self.process_tree:
                #root_processes.append(pid)
                if pid == '0':
                    continue
                # print(self.pslist[node.children[0].pid])
                self.suspect_proc[node.children[0].pid] =self.pslist[node.children[0].pid][2:]+ ', Has a zero root parent'

    # 3. Should find processes running from weired paths

    def malpath(self):
        for pid in self.pslist:
            # print(self.process_path)
            if pid in self.process_path:
                path = self.process_path[pid].lower()

                if not ((path in self.normal_paths) or (path in self.normal_paths_x86)):
                        self.suspect_proc[pid] = self.pslist[pid]

    def find_cmd_child(self, path):

        # Iterate through each row in the CSV file and find the pid of a process executed by cmd.exe
        for pid in self.cmdline:
            # print(pid)
            row = self.cmdline[pid]
            args = row[2]
            # print(args)

            # Find the match in the 'Args' column
            if (path in args) and ('cmd.exe' not in args):
                pid = row[0]
                return pid  # Return the PID if a match is found

        return None  # Return None if no match is found for the given path

    def malcmdline(self):
        for pid in self.cmdline:
            # print(pid)
            row = self.cmdline[pid]
            args = row[2]
            pid = row[0]
                # Find the match in the 'Args' column
            # match = re.search(regex_pattern, args)

            # Iterate through each regex pattern
            for regex_pattern in self.regex_patterns:
                match = re.search(regex_pattern, args)
                if match:
                    path_executed = match.group(0)

                    if pid not in self.suspect_cmdlines:
                        self.suspect_cmdlines[pid] = row
                    #matched_paths.append(path_executed)

                    # print(matched_paths)
                    # When cmd.exe executed, it will create a process. search for this process and add it to the list.
                    # If not found in cmdline list, look for it in psscan list
                    # You should also create the same for powershell.exe, wscript.exe, wmiprvse.exe, rundll32.exe, dllhost.exe, ...

                    if 'cmd\.exe' in regex_pattern:
                        pid = self.find_cmd_child(path_executed)
                        if pid:
                            if pid not in self.suspect_cmdlines:
                            # print(f"Path executed: {path_executed}, PID: {pid}")
                                self.suspect_cmdlines[pid] = self.cmdline[pid]
                        else:
                            print(f"No match found for {path_executed}")

    def malcomm(self, blacklist_file_path):

        # Add option for VPN concentrator

        blacklist_addresses = self.initialize_blacklist_addresses(blacklist_file_path)

        # Define the browser processes you want to exclude
        browsers = ["chrome.exe", "firefox.exe", "iexplore.exe", "edge.exe"]  # Add more if needed

        for pid in self.netscan:

            suspect_comm = []
            if len(self.netscan[pid]) > 1:

                for item in self.netscan[pid]:

                    if 'WinStore.App.e' in item:
                        continue
                    row = item.split(',')

                    if (':::0' in row[2]) or ('0.0.0.0:0' in row[2]) or ('::1:0' in row[2]) :
                        continue

                    else:
                        tmp_addr = row[2].split(':')
                        src_ip = tmp_addr[0]
                        local_port = tmp_addr[1]

                        tmp_addr = row[3].split(':')
                        dst_ip = tmp_addr[0]
                        foreign_port = tmp_addr[1]

                        owner = row[6].lower()  # Convert owner to lowercase for case-insensitive comparison

                         # 1. Any process communicating over port 80, 443, or 8080 that is not a browser
                        if foreign_port in ["80", "443", "8080"] and not any(browser in owner for browser in browsers):
                           suspect_comm.append(','.join(row) + ' A process that is not browser using port: '+foreign_port)

                        # 2. Any browser process not communicating over port 80, 443, or 8080
                        elif any(browser in owner for browser in browsers) and foreign_port not in ["80", "443", "8080"]:
                           suspect_comm.append(','.join(row)+ ' A browser communicating over unusual port: '+ foreign_port)

                        # 3. RDP connections (port 3389), particularly if originating from odd IP addresses. External RDP
                        # connections are typically routed through a VPN concentrator. If the src_ip is not from a VPN concentrator, this is malicious
                        elif foreign_port == "3389" and not src_ip.startswith(("*", "::", "0.0.0.0", "127.0.0.1", "172.16.", "192.168.")):
                           suspect_comm.append(','.join(row)+' External IP communicating directly with RDP port')

                        # 4. Connections to unexplained internal or external IP addresses.
                        # External resources like IP reputation services can also provide additional context.
                        elif dst_ip in blacklist_addresses:
                           suspect_comm.append(','.join(row))

                        # 7. Workstation to workstation connections. Workstations don’t typically RDP, map shares, or authenticate to other workstations.
                        # The expected model is workstations communicate with servers. Workstation to workstation connections often uncover lateral movement.
                        elif ((src_ip.startswith("172.16.") and dst_ip.startswith("172.16.")) or \
                                (src_ip.startswith("192.168.") and dst_ip.startswith("192.168."))):
                           suspect_comm.append(','.join(row)+' Workstation to workstation communication')
                        elif foreign_port in ["5985", "5986"] and not dst_ip.startswith(("0.0.0.0", "127.0.0.1")):
                           suspect_comm.append(','.join(row))

            else:
                row = self.netscan[pid][0].split(',')

                if (':::0' in row[2]) or ('0.0.0.0:0' in row[2]) or ('::1:0' in row[2]) :
                    continue

                else:
                    tmp_addr = row[2].split(':')
                    src_ip = tmp_addr[0]
                    local_port = tmp_addr[1]

                    tmp_addr = row[3].split(':')
                    dst_ip = tmp_addr[0]
                    foreign_port = tmp_addr[1]

                    owner = row[6].lower()  # Convert owner to lowercase for case-insensitive comparison

                     # 1. Any process communicating over port 80, 443, or 8080 that is not a browser
                    if foreign_port in ["80", "443", "8080"] and not any(browser in owner for browser in browsers):
                        suspect_comm.append(','.join(row)+ ' A process that is not browser using port: '+foreign_port)

                    # 2. Any browser process not communicating over port 80, 443, or 8080
                    elif any(browser in owner for browser in browsers) and foreign_port not in ["80", "443", "8080"]:
                        suspect_comm.append(','.join(row)+ ' A browser communicating over unusual port: '+ foreign_port)

                    # 3. RDP connections (port 3389), particularly if originating from odd IP addresses. External RDP
                    # connections are typically routed through a VPN concentrator. If the src_ip is not from a VPN concentrator, this is malicious
                    elif foreign_port == "3389" and not src_ip.startswith(("*", "::", "0.0.0.0", "127.0.0.1", "172.16.", "192.168.")):
                        suspect_comm.append(','.join(row)+' External IP communicating directly with RDP port')

                    # 4. Connections to unexplained internal or external IP addresses.
                    # External resources like IP reputation services can also provide additional context.
                    elif dst_ip in blacklist_addresses:
                        suspect_comm.append(','.join(row))

                    # 7. Workstation to workstation connections. Workstations don’t typically RDP, map shares, or authenticate to other workstations.
                    # The expected model is workstations communicate with servers. Workstation to workstation connections often uncover lateral movement.
                    elif ((src_ip.startswith("172.16.") and dst_ip.startswith("172.16.")) or \
                            (src_ip.startswith("192.168.") and dst_ip.startswith("192.168."))):
                        suspect_comm.append(','.join(row)+' Workstation to workstation communication')
                    elif foreign_port in ["5985", "5986"] and not dst_ip.startswith(("0.0.0.0", "127.0.0.1")):
                        suspect_comm.append(','.join(row))

            # suspect_comm is a temp variable collect suspecious procs for every for loop. At the end of every for loop it add the result to suspect_netscan.
            if suspect_comm:
                if pid not in self.suspect_netscan:
                    self.suspect_netscan[pid] = []
                self.suspect_netscan[pid].append(','.join(suspect_comm))

    def malgetsids(self):

        # suspecious_proc_sids = []
        for pid in self.getsids:
            row = self.getsids[pid][0]
            row = row.split(',')
            process = row[2]
            sid = row[3]

            if sid not in self.normal_sids:
                if process in self.normal_sids.get(sid, []):
                    #print(','.join(row), ' --> Malicious Process: System process running with user account')  # Join the values of the row into a CSV formatted string
                    row.append(' --> Malicious Process: System process running with user account')
                    self.suspecious_proc_sids[pid] = ','.join(row)
                else:
                    self.suspecious_proc_sids[pid] = ','.join(row)
                    #print(','.join(row))  # Join the values of the row into a CSV formatted string
            else:
                # Check if the process is not in the list of normal processes for the SID
                if process not in self.normal_sids[sid]:
                    row.append('Malicious Process: Uknown process running with system account')
                    self.suspecious_proc_sids[pid] =','.join(row)
    def malhandles(self):
        # 1. Identify Anomalous Handles:
        #     Look for handles with unusual or unexpected types, such as file handles for system files or handles to uncommon objects.
        #     Pay attention to handles associated with processes that are not part of the normal system configuration.

        # 2. Examine Handle Counts:
        #     Check for processes with an unusually high number of handles. This could indicate a process that is performing a lot of file or resource operations, which may be suspicious.

        # 3. Investigate Suspicious File Handles:
        #     Focus on file handles and look for processes that have handles to sensitive files, such as system binaries or critical configuration files.
        #     Cross-reference file handles with file paths to determine if any files are accessed from suspicious locations.

        # 4. Analyze Network-Related Handles:
        #     Investigate handles related to network resources, such as sockets or network connections. Look for processes communicating with unusual IP addresses or ports.

        # 5. Correlate with Process Information:
        #     Correlate the handles with other process information like process names, PIDs, and parent-child relationships to identify any processes that appear suspicious or out of place.

        # 6. Check for Code Injection:
        #     Look for processes that have handles to memory sections with Execute (X) permissions. This could indicate code injection or process hollowing, common techniques used by malware.

        # 7. Review Timestamps:
        #     Analyze the timestamps associated with handles. Suspicious processes may have handles with timestamps that don't align with the system's normal behavior.

        # 8. Cross-Reference with Known Malware Indicators:
        #     Compare the handles and associated process information with known indicators of compromise (IOCs) and malware signatures to identify matches.

        # 9. Analyze Parent-Child Relationships:
        #     Look for unusual parent-child relationships between processes. Malware often spawns child processes to perform malicious activities, so identifying such relationships can be crucial.

        # 10. Employ YARA Rules:
        #     Use YARA rules to scan the memory and handles for specific patterns or signatures associated with known malware.

        # 11. Leverage Threat Intelligence:
        #     Consult threat intelligence feeds and databases to check if any of the handles or processes are associated with known malware campaigns or threat actors.

        # 12. Behavioral Analysis:
        #     Consider the overall behavior of processes with suspicious handles. Look for patterns that deviate from normal system behavior, such as excessive file manipulation or network communication.

        # Handles represent “pointers” to objects and it is not  unusual for a process to have a pointer to a file it is reading or writing to. However,
        # this is not the common way code is loaded into a process, so it is unusual to see a DLL file referenced in a process handle (recall that loaded
        # DLLs are typically referenced in the Process Environment Block and displayed via the windows.dlllist plugin).
        # CLASSES Key Registry
        # Storing scripts and data in randomly named registry keys has become a common technique for “fileless” malware, so it pays to keep an eye out strange key

        # The interesting handle in the “msiexec.exe” PID 6192 process is a little more obvious. This process only had 16 file handles. If you think about the
        # purpose of “msiexec.exe” (installing .msi files), finding a reference to “flash_installer.msi” likely indicates what was installed on the system.
        # The name and folder location are of particular interest and ultimately led to recovery of the original infection vector.

        # Named pipes can be identified via “File” handles in Volatility and thus when narrowing your focus, you can filter your output on that handle type
        # (this was not done in the example on the slide since output was filtered with grep for the text “pipe”).

        # However, sometimes things stick out with some connected pipes appending IP addresses or process identifiers as part of the pipe name.

        # Malware commonly sets mutants as a way of "marking" a compromised system to prevent reinfection. The malware will first check for the absence of
        # its mutant before performing an infection. During code analysis, reverse engineers can identify what mutant name is set and checked for by the malware.

        pattern = r'\.dll(?!\.mui)'

        for pid in self.handles:
            for row in self.handles[pid]:
                #print(row)
                if re.search(pattern, row[6]):
                    print(row)

                if '.msi' in row[6]:
                    print(row)

                if 'NamedPipe' in row[6]:
                    print(row)

        print('Here will goes handles function')

    def csv_reader(self, csv_file):
    # Create an empty dictionary to store data with "PID" as the index

        # Read and process the CSV file
        filePath = os.path.join(self.app_path, args.p) + '/'+csv_file

        if csv_file == 'netscan.csv':

            try:
                with open(filePath, 'r', newline='') as file:
                    header = next(file).split()
                    header[2] = ' '.join(header[2:4])
                    header[3] = ' '.join(header[4:6])
                    header[4] = header[6]
                    header[5] = header[7]
                    header[6] = header[8]
                    header[7] = header[9]

                    header = header[0:8]
                    pid_index = header.index('Pid')

                    for line in file:
                        row = line.split()

                        if len(row) < 10:
                            row.insert(4, '')

                        row[7] = ' '.join(row[7:])
                        row = row[0:8]

                        if len(row) > pid_index:
                            pid = row[pid_index]
                            if pid not in self.netscan:
                                self.netscan[pid] = []
                        self.netscan[pid].append(','.join(row))

            except FileNotFoundError:
                print(f"File '{csv_file}' not found.")
        else:

            try:
                with open(filePath, 'r', newline='') as csvfile:
                    csvreader = csv.reader(csvfile)
                    header = next(csvreader)  # Read the header

                    # Find the index of the "PID" column
                    pid_index = header.index('PID')

                    path_index = 0
                    if csv_file == 'dlllist.csv':
                        path_index = header.index('Path')

                    if csv_file == 'pslist.csv':
                    # Store the row as a comma-separated string in the dictionary
                        for row in csvreader:
                            if len(row) > pid_index:
                                pid = row[pid_index]
                            self.pslist[pid] = ','.join(row)

                    elif csv_file == 'psscan.csv':
                        for row in csvreader:
                            if len(row) > pid_index:
                                pid = row[pid_index]
                            self.psscan[pid] = ','.join(row)

                    elif csv_file == 'cmdline.csv':
                        for row in csvreader:

                            if len(row) >= 2:
                                if 'process exited' in row[3]:
                                    continue
                                pid = row[pid_index]
                                self.cmdline[pid] = row[1:]

                    elif csv_file == 'dlllist.csv':
                        for row in csvreader:
                            if len(row) > pid_index:
                                pid = row[pid_index]
                                if pid not in self.dlllist:
                                    self.dlllist[pid] = []       # Every process PID has more than one DLL. Keep all dlls for every process in the same pid
                            self.dlllist[pid].append(','.join(row))

                            if len(row) > path_index:
                                path = row[path_index]

                                if '.exe' in path:       #Skipp process path
                                    self.process_path[pid] = path
                                    continue
                                if path not in self.dll_stacking:
                                    self.dll_stacking[path] = 0
                            self.dll_stacking[path] += 1

                    elif csv_file == 'getsids.csv':
                        for row in csvreader:
                            if len(row) == 0:
                                break

                            if len(row) > pid_index:
                                pid = row[pid_index]
                                if pid not in self.getsids:
                                    self.getsids[pid] = []
                            self.getsids[pid].append(','.join(row))

                    elif csv_file == 'handles.csv':
                        filter = ['File', 'Key', 'Mutant']  #

                        for row in csvreader:
                            #print(','.join(row))
                            if  len(row) == 0:
                                continue
                            if row[5] not in filter or row[7] == '':
                                continue

                            if len(row) > pid_index:
                                pid = row[pid_index]
                                if pid not in self.handles:
                                    self.handles[pid] = []       # Every process PID has a lot of handles. Keep all handles for every process in the same pid
                            self.handles[pid].append(row[1:])

            except FileNotFoundError:
                print(f"File '{csv_file}' not found.")

    def collect_evidence(self, pid):
        """This funtion collects artifacts from all memory objects. This how you can print it print(self.artifact_bag['4'][0]['pslist'])"""
        #print(self.artifact_bag['4'][0]['pslist'])
        #print("I will collect artifacts using this funtion!")
        #print('This is the blueprint for this bag: ioc_collection = [pid: {<plugin1>: <row from plugin csv file>}, {<plugin2>: <row from plugin2 csv file>'+'}]')
        #self.evidence_bag[4] = [{'pslist':'0,2876,868,WmiPrvSE.exe,0x8c88b14e7580,10,-,0,False,2018-08-30 13:52:26.000000 ,N/A,Disabled'}]

        # if pid not in self.evidence_bag:
        #     self.evidence_bag[pid] = []

        # Collect evidence from PsList

        # reason = 'High privilege account SID: '

        # if pid in self.getsids:

        #     if len(self.getsids[pid]) > 1:

        #         for item in self.getsids[pid]:

        #             if ('Everyone' in item) or ('Users' in item) or ('This Organization' in item):
        #                 continue

        #             columns = item.split(',')

        #             if 'Local System' == columns[4]:
        #                 reason += ' '.join(columns[3:])

        #             # Don't forget to add the enrichment from malgetsids function.
        #             elif '544' in columns[3]:
        #                 reason += 'The process spawned from a user context with Administrator privileges: '+' '.join(columns[3:])

        #             if pid in self.suspecious_proc_sids:
        #                 reason2 = self.suspecious_proc_sids[pid].split(',')
        #                 if len(reason2) == 6:
        #                     reason += reason2[5]
        ppid = ''
        if pid in self.pslist:

            #print(self.psscan[pid])
            # evidence = {'pslist': self.pslist[pid]}
            #self.evidence_bag[pid].append(evidence)
            row = []
            row.append('pslist')
            columns = self.pslist[pid].split(',')
            row.append(columns[1])

            ppid = columns[2]
            row.append(columns[2])
            row.append(columns[3])

            path = ''
            if pid in self.process_path:
                path = self.process_path[pid]

            # You need to change this list because it is not clean, missing a lot of system files and dlls.
            # if path in self.whitelist_paths and self.score[pid] < 40:
            #     return

            row.append(path)
            row.append(columns[9])

            if pid in self.suspect_proc:
                reason = self.suspect_proc[pid].split(',')[11]
                row.append(reason)

            self.evidence_bag.append(','.join(row))
        # print(pid, ': ',self.evidence_bag)

        # Collect evidence from PsScan
        if (pid in self.psscan) and (pid not in self.pslist):
            #print(self.psscan[pid])
            # evidence = {'psscan': self.psscan[pid]}
            # self.evidence_bag[pid].append(evidence)

            row = []
            row.append('psscan')
            columns = self.psscan[pid].split(',')
            row.append(columns[1])
            row.append(columns[2])
            row.append(columns[3])

            path = ''
            if pid in self.process_path:
                path = self.process_path[pid]

            row.append(path)
            row.append(columns[9])

            # 4. Should hidden/Unlinked processes like processes found in psscan, but not in pslist
            # Hidden processes: Also add more details such as number of threads, parent process, ....
            if row[1] not in self.pslist:
                row.append('Hidden/Unlinked: The process found in psscan, but not in pslist')
            self.evidence_bag.append(','.join(row))

        # Collect evidence from DllList

        if pid in self.dlllist:

            if len(self.dlllist[pid]) > 1:

                for item in self.dlllist[pid]:

                    row =[]
                    row.append('dlllist')
                    # evidence = {'dlllist': self.dlllist[pid]}
                    # self.evidence_bag[pid].append(evidence)

                    columns = item.split(',')

                    if '.exe' in columns[6]:
                        continue

                    # We can add option to use this filtering or not. If we want to find malicious dlls we can use filtering, but
                    # if we want o understand the malcious process capabilities, we ignore filtering becuase used libraries can tell
                    # you the malware can do

                    if columns[6].lower() in self.whitelist_paths:
                        continue

                    if self.dll_stacking[columns[6]] > 2:       # Only keep dlls with low frequency as malicious dlls always rare
                        continue

                    row.append(columns[1])
                    row.append(ppid)
                    row.append(columns[2])
                    row.append(columns[6])
                    row.append(columns[7])
                    row.append('Whitelist the clean DLLs')
                    self.evidence_bag.append(','.join(row))
            else:
                row =[]
                row.append('dlllist')
                # evidence = {'dlllist': self.dlllist[pid]}
                # self.evidence_bag[pid].append(evidence)

                columns = self.dlllist[pid].split(',')
                if self.dll_stacking[columns[6]] > 5:       # Only keep dlls with low frequency as malicious dlls always rare
                    return
                row.append(columns[1])
                row.append(ppid)
                row.append(columns[2])
                row.append(columns[6])
                row.append(columns[7])
                row.append('Whitelist the clean DLLs')
                self.evidence_bag.append(','.join(row))

        # Collect evidence from CmdLine
        if pid in self.cmdline:
            # evidence = {'cmdline': self.cmdline[pid]}
            # self.evidence_bag[pid].append(evidence)

            row =[]
            row.append('cmdline')
            #columns = self.cmdline[pid].split(',')
            row.append(self.cmdline[pid][0])          # Add pid to the list
            row.append(ppid)
            row.append(self.cmdline[pid][1])                    # Add process name

            pattern = r'"?([C|c]:\\[a-zA-Z0-9%\\ \(\)\.]*\.[eE][xX][eE]|[%\\]SystemRoot[%\\][a-zA-Z0-9%\\]*\.exe)\b'
            paths = re.findall(pattern, self.cmdline[pid][2])

            if paths:
                paths = paths[0].strip('"')

            row.append(path)
            row.append('')
            row.append(self.cmdline[pid][2].replace('"', ''))
            self.evidence_bag.append(','.join(row))

        if pid in self.netscan:

            if len(self.netscan[pid]) > 1:

                for item in self.netscan[pid]:
                    row =[]
                    row.append('netscan')
                    # evidence = {'dlllist': self.dlllist[pid]}
                    # self.evidence_bag[pid].append(evidence)

                    columns = item.split(',')
                    row.append(columns[5])
                    row.append(ppid)
                    row.append(columns[6])
                    row.append(columns[1]+' ' + ' --> '.join(columns[2:4])+' '+columns[4])
                    row.append(columns[7])
                    row.append('I Will add info later')
                    self.evidence_bag.append(','.join(row))
            else:
                # evidence = {'netscan': self.netscan[pid]}
                # self.evidence_bag[pid].append(evidence)

                row =[]
                row.append('netscan')
                # evidence = {'dlllist': self.dlllist[pid]}
                # self.evidence_bag[pid].append(evidence)

                columns = self.netscan[pid][0].split(',')
                row.append(columns[5])
                row.append(ppid)
                row.append(columns[6])
                row.append(columns[1]+' ' + ' --> '.join(columns[2:4])+' '+columns[4])
                row.append(columns[7])
                row.append('I Will add info later')

                self.evidence_bag.append(','.join(row))

        if pid in self.getsids:

            if len(self.getsids[pid]) > 1:

                local_system = False
                for item in self.getsids[pid]:

                    if any(keyword in item for keyword in ['Power', 'Logon Session', 'Authentication Authority Asserted Identity', 'PlugPlay', 'Users',
                                                           'This Organization', 'Everyone', 'DcomLaunch', 'High Mandatory Level']):
                        continue

                    row =[]
                    row.append('getsids')
                    columns = item.split(',')

                    if columns[4] == '-':
                        continue
                    row.append(columns[1])
                    row.append(ppid)
                    row.append(columns[2])
                    row.append(columns[3])
                    row.append('')

                    # Don't forget to add the enrichment from malgetsids function.

                    reason = columns[4]
                    if '544' in columns[3]:
                        reason = 'The process spawned from a user context with high privileges account: ' + reason

                    if pid in self.suspecious_proc_sids:
                        reason2 = self.suspecious_proc_sids[pid].split(',')
                        if len(reason2) == 6:
                            reason = reason2[5] +': '+ reason

                    row.append(reason)
                    self.evidence_bag.append(','.join(row))

    def create_SuperMemoryParser(self):

        header = ['Source Name', 'PID', 'Process Name', 'Path', 'Timestamps', 'Long Description']

        with open('SuperMemoryParser.csv', 'w') as mparser:
            mparser.write('\n'.join(self.evidence_bag))

    def malicious_weight(self):

        if self.pslist:
            for pid in self.pslist:
                weight = 0
                if pid in self.suspect_proc:
                    # if pid =='9036':
                    #     print(f"Found {pid} in suspect_proc: {self.suspect_proc[pid]}")
                    weight += 10

                    if 'Has a zero root parent' in self.suspect_proc[pid]:
                        weight += 30

                if pid in self.suspecious_proc_sids:
                    # if pid =='9036':
                    #     print(f"Found {pid} in suspecious_proc_sids: {self.suspecious_proc_sids[pid]}")
                    weight += 10

                if pid in self.suspect_cmdlines:
                    # if '9036' in self.suspect_cmdlines[pid]:
                    #     print(f"Found {pid} in suspect_cmdlines: {self.suspect_cmdlines[pid]}")
                    weight += 10

                if pid in self.suspect_netscan:
                    # if pid =='9036':
                    #     print(f"Found {pid} in suspect_netscan: {self.suspect_netscan[pid]}")
                    weight += 10

                # print(self.baseline_proc['1096'])
                if pid in self.baseline_proc:
                    # if pid =='9036':
                    #     print(f"Found {pid} in baseline_proc: {self.baseline_proc[pid]}")
                    weight += 10

                if weight > 30:
                    self.score[pid] = weight

# Function to add a child node to a parent node
def add_child(parent, child):
    parent.children.append(child)

# Function to recursively print the tree with '*' depth

def print_process_tree(node, depth=0):

    #     # Print the node's information with depth represented by '*' characters
    #     print("*" * depth + ' '*(12-depth)+ f"{node.pid}, {node.ppid}, {node.image_file_name}")
    #     # Recursively print children

    if node:
            # Print the header only for the root node
        if depth == 0:
            print(f"{'':<12}{'PID':<12}{'PPID':<12}{'Image File Name'}")
        # Check if the node is not the root (0, 0)
        if node.pid != '0':
            # Print the node's information with depth represented by '*' characters
            print("*" * depth + ' ' * (12 - depth) + f"{node.pid:<12}{node.ppid:<12}{node.image_file_name}")
        # Recursively print children
        for child in node.children:
            print_process_tree(child, depth + 1)

def get_application_path():
    try:
        if getattr(sys, 'frozen', False):
            application_path = os.path.dirname(os.path.realpath(sys.executable))
        else:
            application_path = os.path.dirname(os.path.realpath(__file__))
        if "~" in application_path and os_platform == "windows":
            # print "Trying to translate"
            # print application_path
            application_path = win32api.GetLongPathName(application_path)
        #if args.debug:
        #    logger.log("DEBUG", "Init", "Application Path: %s" % application_path)
        return application_path
    except Exception as e:
        print("Error while evaluation of application path")
        traceback.print_exc()
        if args.debug:
            sys.exit(1)

def main(memory):

    """
    Argument parsing function
    :return:
    """
    # for pid in memory.score:
    #     print(pid, ': ', memory.score[pid])

    # child_pid = '8260'  # Replace with the PID of the child you want to find the parent for
    # # parent_node = memory.get_parent(child_pid)

    # memory.find_parent_recursive(child_pid)
    # parent_node = memory.get_parent(child_pid)

    # while parent_node.image_file_name != 'svchost.exe':

    #     parent_node = memory.get_parent(child_pid)

    #     if parent_node:
    #         print(f"Parent PID: {parent_node.pid}, Image File Name: {parent_node.image_file_name}")
    #     else:
    #         print(f"No parent found for PID {child_pid}.")

    #     child_pid = parent_node.pid

    pid_list = [key for key in memory.score]

    for pid in pid_list:
        memory.collect_evidence(pid)
        #print(pid, ': ', memory.score[pid])
        # parent = memory.pslist[pid].split(',')[2]

        memory.find_parent_recursive(pid)

        # if parent not in memory.score:
        #     memory.collect_evidence(parent)

    memory.create_SuperMemoryParser()

    while True:

        print("Available plugins to explore memory data:")
        for index, command in enumerate(memory.plugins, start=1):
            print(f"{index}. {command}")

        print("0. Exit")

        try:
            choice = int(input("Enter the plugin number you want to select (or 0 to quit): "))
            if choice > len(memory.plugins):
                print("\nInvalid input. Please enter a valid number.")
                continue
            if choice == 0:
                print(memory.score)
                break
            selected_command = memory.plugins[choice - 1]
            corresponding_file = memory.csv_files[choice - 1]
            print(f"Selected command: {selected_command}")
            print(f"Corresponding CSV file: {corresponding_file}")

            if selected_command == 'pslist':
                for pid in memory.pslist:
                    print(memory.pslist[pid])

            if selected_command == 'psscan':
                for pid in memory.psscan:
                    print(memory.psscan[pid])

            if selected_command == 'pstree':
                print_process_tree(memory.process_tree['0'])

            if selected_command == 'dlllist':
                for pid in memory.dlllist:
                    print(memory.dlllist[pid])

            if selected_command == 'cmdline':
                for pid in memory.cmdline:
                    print(memory.cmdline[pid])

            if selected_command == 'netscan':
                for pid in memory.netscan:
                    print(memory.netscan[pid])

            if selected_command == 'getsids':
                for pid in memory.getsids:
                    print(memory.getsids[pid])

            print("1. Display only suspecious processes")
            print("2. Continue with full list")
            print("0. Back")

            try:
                choice = int(input("Select one option (or 0 go back): "))
                if choice == 0:
                    continue

                elif choice == 1:
                    while True:
                        print("Display only suspecious processes:")
                        for index, command in enumerate(memory.analyze, start=1):
                            print(f"{index}. {command}")
                        print("0. Back")

                        try:
                            choice = int(input("Enter the number corresponding to the analysis you want to select (or 0 go back): "))

                            if choice == 0:
                                break

                            elif choice <= len(memory.analyze):
                                if choice == 1:
                                    #suspect_proc = memory.malproc(memory.process_tree['4'])
                                    for pid in memory.suspect_proc:
                                        print(memory.suspect_proc[pid])
                                    print('\n\nThe above processes are more likely the most suspecious processes. Select one of them by PID\n\n')

                                elif choice == 3:
                                    for pid in memory.suspect_cmdlines:
                                        print(memory.suspect_cmdlines[pid])

                                elif choice == 5:
                                    print('\nThis function have not been implemted, please do it!')
                                    continue

                                    suspect_proc = memory.malcomm('blacklist.txt')
                                    print('\n'.join(suspect_proc))

                                elif choice == 6:
                                    for pid in memory.suspect_netscan:
                                        print(memory.suspect_netscan[pid])

                                elif choice == 7:
                                    for pid in memory.baseline_proc:
                                        print(memory.baseline_proc[pid].replace('"', ''))

                                elif choice == 8:
                                    memory.malgetsids()

                                    for pid in memory.suspecious_proc_sids:
                                        print(memory.suspecious_proc_sids[pid])

                                while True:
                                    # Prompt for the PID to search for
                                    pid_to_search = input("\n\nEnter the suspecious process PID that you want to collect as evidence (or type 'back' to go back to plugins selection): ")

                                    if pid_to_search.lower() == 'back':
                                        break
                                    memory.collect_evidence(pid_to_search)
                                            #print(memory.pslist[pid_to_search])

                            else:
                                print("Invalid choice. Please select a valid number.")

                        except ValueError:
                            print("Invalid input. Please enter a valid number.")

                elif choice == 2:
                    while True:
                        # Prompt for the PID to search for
                        pid_to_search = input("\n\nEnter the suspecious process PID that you want to collect as evidence (or type 'back' to go back to plugins selection): ")

                        if pid_to_search.lower() == 'back':
                            break
                        memory.collect_evidence(pid_to_search)
                                #print(memory.pslist[pid_to_search])
                        #print(memory.evidence_bag)
                else:
                    print("Invalid choice. Please select a valid number.")

            except ValueError:
                print("Invalid input. Please enter a valid number.")

            # for key, value in memory.evidence_bag.items():
            #     formatted_output = []
            #     for item in value:
            #         for entry_type, entry_data in item.items():
            #             if entry_type == 'dlllist':
            #                 formatted_output.append(f'{entry_type}:')
            #                 formatted_output.extend([f'    {line}' for line in entry_data])

            #             elif entry_type == 'netscan':
            #                 formatted_output.append(f'{entry_type}:')
            #                 formatted_output.extend([f'    {line}' for line in entry_data])
            #             else:
            #                 formatted_output.append(f'{entry_type}: {entry_data}')
            #     print("\n".join(formatted_output))
            #     print('\n')

        except ValueError:
            print("Invalid input. Please enter a valid number.")

if __name__ == "__main__":

    # Argument parsing

    # Parse Arguments
    parser = argparse.ArgumentParser(description='AutVol3 - Simple Memoray Image Analyzer')
    parser.add_argument('-p', help='Path to memory image and csv files like pslist, netscan ...', metavar='path', default='')
    parser.add_argument('--version', action='store_true', help='Shows welcome text and version of AutoVol3, then exit', default=False)

    # Add option for baseline
    # Add option for image profile if using vol2
    # Add option for blacklist IPs
    # Add option for chosing memory image file
    # Add option to add VPN concentrator

    args = parser.parse_args()

    if not args.p:
        print('Must specify memory image path!')
        parser.print_help()
        sys.exit(1)

    autovol3 = AutoVol3(args)

    # Show version
    if args.version:
        sys.exit(0)
    #suspect_proc = autovol3.malcomm('blacklist.txt')
    #autovol3.malgetsids()
    main(autovol3)

    # print('This is the last push')


import numpy as np
import matplotlib.pyplot as plt

black1 = np.zeros((50,50))
black2 = black1.copy()
white1 = np.ones((50,50))
white2 = white1.copy()
row1 = np.hstack((black1,white1))
row2 = np.hstack((white2, black2))
full = np.vstack((row1,row2))
plt.figure()
plt.imshow(full, cmap='gray')
plt.show()


import random
import sqlite3

def sql_create():
    global db, cursor
    db = sqlite3.connect("bot.sqlite3")
    cursor = db.cursor()

    if db:
        print("База данных подключена!")

    db.execute("CREATE TABLE IF NOT EXISTS anketa "
               "(id INTEGER PRIMARY KEY, "
               "name VARCHAR (100), "
               "name VARCHAR (100), "
               "direction VARCHAR(100), "
               "age INTEGER (20), "
               "group VARCHAR(10), ")
    db.commit()

async def sql_command_insert(state):
    async with state.proxy() as data:
        cursor.execute("INSERT INTO anketa VALUES (?, ?, ?, ?, ?, ?, ?)",
                       tuple(data.values()))
        db.commit()

async def sql_command_random():
    result = cursor.execute("SELECT * FROM anketa").fetchall()
    random_user = random.choice(result)
    return random_user

async def sql_command_all():
    return cursor.execute("SELECT * FROM anketa").fetchall()

async def sql_command_delete(id):
    cursor.execute("DELETE FROM anketa WHERE id = ?", (id,))
    db.commit()


import cv2
from pykuwahara import kuwahara

image = cv2.imread('selfie.jpg')
image = (image / 255).astype('float32')     # pykuwahara supports float32 as well

lab_image = cv2.cvtColor(image, cv2.COLOR_BGR2Lab)
l, a, b = cv2.split(lab_image)

hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
h, s, v = cv2.split(hsv_image)

filt1 = kuwahara(image, method='gaussian', radius=5, sigma=2., image_2d=l)
filt2 = kuwahara(image, method='gaussian', radius=5, sigma=2., image_2d=v)
#filt = kuwahara(lab_image, method='gaussian', radius=5, sigma=2., image_2d=l)

cv2.imwrite('selfie-kfilt-gaus1.jpg', filt1 * 255)
cv2.imwrite('selfie-kfilt-gaus2.jpg', filt2 * 255)
#cv2.imwrite('selfie-kfilt-gaus.jpg', cv2.cvtColor(filt, cv2.COLOR_Lab2BGR) * 255)


# Support Vector Machine (SVM)
# from https://www.superdatascience.com/machine-learning/

# Part 1 - Data Preprocessing

# Importing the libraries
import os
import glob
import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt

train_dir = os.path.join(os.path.curdir, "data/train")
validation_dir = os.path.join(os.path.curdir, "data/validate")

train_dirs = []
train_dirs.append(os.path.join(train_dir, "cima"))
train_dirs.append(os.path.join(train_dir, "baixo"))
train_dirs.append(os.path.join(train_dir, "esq"))
train_dirs.append(os.path.join(train_dir, "dir"))

validation_dirs = []
validation_dirs.append(os.path.join(validation_dir, "cima"))
validation_dirs.append(os.path.join(validation_dir, "baixo"))
validation_dirs.append(os.path.join(validation_dir, "esq"))
validation_dirs.append(os.path.join(validation_dir, "dir"))

# Preparando os dados de treinamento
X = []
y = []
for dir in train_dirs:
    classe = dir.rsplit('/')[-1]
    if classe == 'cima':
        classe_num = 0
    if classe == 'baixo':
        classe_num = 1
    if classe == 'esq':
        classe_num = 2
    if classe == 'dir':
        classe_num = 3

    files = os.listdir(dir)
    for file in files:
        image = cv.imread(dir+"/"+file, 0)
        # flatten transforma o array 2D de imagem em 1D
        X.append(image.flatten())
        y.append(classe_num)

X_train = np.array(X)
y_train = np.array(y)
print(X_train.shape)
print(y_train.shape)

# Preparando os dados de validação
X = []
y = []
for dir in validation_dirs:
    classe = dir.rsplit('/')[-1]
    if classe == 'cima':
        classe_num = 0
    if classe == 'baixo':
        classe_num = 1
    if classe == 'esq':
        classe_num = 2
    if classe == 'dir':
        classe_num = 3

    files = os.listdir(dir)
    for file in files:
        image = cv.imread(dir+"/"+file, 0)
        # flatten transforma o array 2D de imagem em 1D
        X.append(image.flatten())
        y.append(classe_num)

X_val = np.array(X)
y_val = np.array(y)
print(X_val.shape)
print(y_val.shape)

# Feature Scaling
from sklearn.preprocessing import StandardScaler
sc = StandardScaler()
X_train = sc.fit_transform(X_train)
X_val = sc.transform(X_val)

# Fitting SVM to the Training set
from sklearn.svm import SVC
from sklearn.pipeline import Pipeline

#classifier = SVC(kernel = 'linear', random_state = 0)
classifier = Pipeline((
("scaler", StandardScaler()),
("svm_clf", SVC(kernel="rbf", gamma=0.01, C=1000000))
))

classifier.fit(X_train, y_train)

# Predicting the Test set results
y_pred = classifier.predict(X_val)

print(y_val[0:35])
print(y_pred[0:35])

# Making the Confusion Matrix
from sklearn.metrics import confusion_matrix
cm = confusion_matrix(y_val, y_pred)

# matriz de confusão
print(cm)

print("Score da validação: ", classifier.score(X_val, y_val))

# Salvando modelo
import joblib as jl

model_file = "svm_eyes.pkl"
jl.dump(classifier, model_file)


# Testzone para o apêndice

import apendice as ap
import matplotlib.pyplot as plt

#%% Plotando os sharp_test

## Sinal 1

vector1 = ap.sharpe_test_1(combined = combined, tam_p = 0.1, qtde = 200, inverso = True)

plt.hist(vector1.iloc[:,0], bins = 40)
plt.ylabel('Quantas vezes')
plt.xlabel('Sharpe')
plt.title('Para o Sinal 1')
plt.show()

## Sinal 2

vector2 = ap.sharpe_test_2(combined = combined, tam_p = 0.1, qtde = 200, inverso = True)

plt.hist(vector2.iloc[:,0], bins = 40)
plt.ylabel('Quantas vezes')
plt.xlabel('Sharpe')
plt.title('Para o Sinal 2')
plt.show()

## Sinal 3

vector3 = ap.sharpe_test_3(combined = combined, tam_p = 0.1, qtde = 200, inverso = True, dias = 17)

plt.hist(vector3.iloc[:,0], bins = 40)
plt.ylabel('Quantas vezes')
plt.xlabel('Sharpe')
plt.title('Para o Sinal 3')
plt.show()

## Sinal 4

vector4 = ap.sharpe_test_4(combined = combined, tam_p = 0.1, qtde = 200, inverso = True, dias = 21)

plt.hist(vector4.iloc[:,0], bins = 40)
plt.ylabel('Quantas vezes')
plt.xlabel('Sharpe')
plt.title('Para o Sinal 4')
plt.show()

#%% Avaliando a normalização via correlação média
## Se quer ver a avaliação, insira 1

ap.av_corr_med(insert = 0)


# -*- coding: utf-8 -*-
"""
Created on Fri Aug 17 05:15:30 2018
@author: LeeMH
"""
from bs4 import BeautifulSoup
from datetime import datetime, timedelta
import requests
import re

from utils.processor_checker import timeit
from gobble.module.crawler import Crawler

class NaverRealtimeCrawler(Crawler):
    '''
    네이버 증권 실시간 속보 크롤러
    '''
    def __init__(self):
        super().__init__()

    # func (기능): all (오늘 날짜의 모든 뉴스를 크롤링), new (1page만 크롤링)
    @timeit
    def create_url_list(self, func):
        req = self.request_get(self.real_time_list.format(self.rt_today, 1), self.user_agent)
        print(self.real_time_list.format(self.rt_today, 1))
        soup = self.html_parser(req)
        url_list = []
        if func == 'new':
            url_data = self.find_navernews_url(soup, url_list)

        elif func == 'all':
            pgRR = self.soup_find(soup, 'td', {'class':'pgRR'})
            last_page = self.soup_find(pgRR, 'a')['href'][-3:]
            last_page = re.findall("\d+", last_page)[0]

            sub_list = []
            for i in range(1,int(last_page)+1):
                req = self.request_get(self.real_time_list.format(self.rt_today, i), self.user_agent)
                sub_soup = self.html_parser(req)
                url_list += self.find_navernews_url(sub_soup, sub_list)
            url_data = url_list
        else:
            print("Choose between 'all' and 'new'")
        url_data = list(set(url_data))
        print(len(url_data))
        real_time_url = [self.soup_find(sub, 'a')['href'].replace('§', '&sect') for sub in url_data]
        return real_time_url

    @timeit
    def get_data(self, url_list, checklist):
        url_list = url_list
        data_list = []
        for url in url_list:
            self.req = self.request_get(self.fin_nhn.format(url), self.user_agent)
            soup = self.html_parser(self.req)
            url = self.fin_nhn.format(url)
            title = self.soup_find(soup, 'h3').text.replace('\n','').replace('\t','').strip()
            checkurl = re.sub(r'&page=\d+$', '', url)
            if checkurl in checklist:
                print('Already up-to-date.')
                continue
            upload_time = self.soup_find(soup, 'span', {'class':'article_date'}).text
            media = self.soup_find(soup, 'span', {'class':'press'}).img['title']
            data_dict = {'title':title, 'media':media, 'url':checkurl, 'data_type':'R', 'upload_time': upload_time}
            data_list.append(data_dict)
        print(len(data_list))
        return data_list


from wsgiref.simple_server import make_server
import psutil,datetime
import sqlite3

def interview_scores(environ, start_response):
    conn_result=sqlite3.connect("results.sqlite")
    cursor_result = conn_result.cursor()
    print_results = cursor_result.execute("select * from interview_results")
    status = '200 OK'
    headers = [('Content-type', 'html; charset=utf-8')]
    start_response(status, headers)
    message = "
Interview Scores
"
    message +="
"
    message += "
"
    message +="Username	"
    message +="Correct Answers	"
    message +="Total Questions	"
    message +="Result
"

    for data in print_results:
        message += "
"
        message += ""+(str(data[0]))+"	"
        message += ""+(str(data[1]))+"	"
        message += ""+(str(data[2]))+"	"
        message += ""+(str(data[3]))+"	"
        message += "
"
    return[bytes(message,'utf-8')]

httpd = make_server('', 8000,interview_scores)
print("Serving on port 8000...")

httpd.serve_forever()


#!/usr/bin/python3

import os
import argparse

BASE_CMD = 'gtimeout 60 stack exec -- hplus \"{query}\" --stop-refine\
            --stop-threshold=10 --cnt=5'

def main():
    parser = argparse.ArgumentParser(description='Run a single hoogle+ query')
    parser.add_argument('query', help='the signature to be searched')
    args = parser.parse_args()
    os.system(BASE_CMD.format(query=args.query))

if __name__ == "__main__":
    main()


#!/bin/python3
import os
import pandas
import argparse
import traceback
import numpy as np
from datetime import datetime, timedelta

# lisflood
import lisf1
from lisflood.global_modules.decorators import Cache

from liscal import hydro_model, templates, config, subcatchment, calibration, objective

class ScalingModel():
    def __init__(self, cfg, subcatch, lis_template, lock_mgr, objective):

        self.cfg = cfg
        self.subcatch = subcatch

        self.lis_template = lis_template

        self.lock_mgr = lock_mgr

        self.objective = objective

        self.start = cfg.forcing_start.strftime('%d/%m/%Y %H:%M')
        self.end = cfg.forcing_end.strftime('%d/%m/%Y %H:%M')

    def init_run(self):

        # dummy Individual, doesn't matter here
        param_ranges = self.cfg.param_ranges
        Individual = 0.5*np.ones(len(param_ranges))

        cfg = self.cfg

        gen = self.lock_mgr.get_gen()
        run_id = str(gen)

        out_dir = os.path.join(self.subcatch.path_out, run_id)
        os.makedirs(out_dir, exist_ok=True)

        parameters = self.objective.get_parameters(Individual)

        prerun_file, run_file = self.lis_template.write_init(run_id, self.start, self.end, self.start, self.end, cfg.param_ranges, parameters)

        # -i option to exit after initialisation, we just load the inputs map in memory
        lisf1.main(prerun_file, '-i')
        lisf1.main(run_file, '-i')

        # store lisflood cache size to make sure we don't load anything else after that
        self.lisflood_cache_size = Cache.size()

    def run(self, Individual):

        cfg = self.cfg

        gen = self.lock_mgr.get_gen()
        run = self.lock_mgr.increment_run()
        print('Generation {}, run {}'.format(gen, run))

        run_id = '{}_{}'.format(gen, run)

        out_dir = os.path.join(self.subcatch.path_out, run_id)
        os.makedirs(out_dir, exist_ok=True)

        parameters = self.objective.get_parameters(Individual)

        prerun_file, run_file = self.lis_template.write_template(run_id, self.start, self.end, self.start, self.end, cfg.param_ranges, parameters)

        lisf1.main(prerun_file, '-v')
        lisf1.main(run_file, '-v')

def scaling_subcatchment(cfg, obsid, subcatch, n_runs):

    lis_template = templates.LisfloodSettingsTemplate(cfg, subcatch)

    lock_mgr = calibration.LockManager(cfg.num_cpus)

    lock_mgr.set_gen(cfg.num_cpus)

    obj = objective.ObjectiveKGE(cfg, subcatch, read_observations=False)

    model = ScalingModel(cfg, subcatch, lis_template, lock_mgr, obj)

    # load forcings and input maps in cache
    # required in front of processing pool
    # otherwise each child will reload the maps
    model.init_run()

    dummy_params = [0.5*np.ones(len(cfg.param_ranges)) for i in range(n_runs)]

    scaling_map, pool = lock_mgr.create_mapping()

    mapped = list(scaling_map(model.run, dummy_params))

    if pool:
        pool.close()

class ConfigScaling(config.Config):

    def __init__(self, settings_file, num_cpus):
        super().__init__(settings_file)

        self.num_cpus = int(num_cpus)

        # paths
        self.subcatchment_path = self.parser.get('Path','subcatchment_path')

        # Date parameters
        self.forcing_start = datetime.strptime(self.parser.get('Main','forcing_start'),"%d/%m/%Y %H:%M")  # Start of forcing
        self.forcing_end = datetime.strptime(self.parser.get('Main','forcing_end'),"%d/%m/%Y %H:%M")  # Start of forcing
        self.timestep = int(self.parser.get('Main', 'timestep'))  # in minutes
        self.prerun_timestep = 360  # in minutes
        if self.prerun_timestep != 360 and self.prerun_timestep != 1440:
            raise Exception('Pre-run timestep {} not supported'.format(self.prerun_timestep))

        # Load param ranges file
        self.param_ranges = pandas.read_csv(self.parser.get('Path','param_ranges'), sep=",", index_col=0)

        # template
        self.lisflood_template = self.parser.get('Templates','LISFLOODSettings')

        # pcraster commands
        self.pcraster_cmd = {}
        for execname in ["pcrcalc", "map2asc", "asc2map", "col2map", "map2col", "mapattr", "resample", "readmap"]:
            self.pcraster_cmd[execname] = execname

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('settings', help='Settings file')
    parser.add_argument('stations_data', help='Stations metadata CSV file')
    parser.add_argument('obsid', help='Station obsid')
    parser.add_argument('n_cpus', help='Station obsid')
    parser.add_argument('n_runs', help='Station obsid')
    args = parser.parse_args()

    print('  - obsid: {}'.format(args.obsid))
    print('  - settings file: {}'.format(args.settings))
    obsid = int(args.obsid)
    cfg = ConfigScaling(args.settings, args.n_cpus)

    print(">> Reading stations.csv file...")
    stations = pandas.read_csv(args.stations_data, sep=",", index_col=0)
    try:
        station_data = stations.loc[obsid]
    except KeyError as e:
        print(stations)
        raise Exception('Station {} not found in stations file'.format(obsid))

    subcatch = subcatchment.SubCatchment(cfg, obsid, station_data, create_links=False)

    scaling_subcatchment(cfg, obsid, subcatch, int(args.n_runs))


# This file is Copyright 2020 Volatility Foundation and licensed under the Volatility Software License 1.0
# which is available at https://www.volatilityfoundation.org/license/vsl-v1.0

import json
import logging
import ntpath
import os
import re
from typing import List, Dict, Union

from volatility3.framework import (
    renderers,
    interfaces,
    objects,
    exceptions,
    constants,
    layers,
)
from volatility3.framework.configuration import requirements
from volatility3.framework.renderers import format_hints
from volatility3.framework.symbols.windows.extensions import registry
from volatility3.plugins.windows import pslist
from volatility3.plugins.windows.registry import hivelist

vollog = logging.getLogger(__name__)

def find_sid_re(
    sid_string, sid_re_list
) -> Union[str, interfaces.renderers.BaseAbsentValue]:
    for reg, name in sid_re_list:
        if reg.search(sid_string):
            return name
    return renderers.NotAvailableValue()

class GetSIDs(interfaces.plugins.PluginInterface):
    """Print the SIDs owning each process"""

    _version = (1, 0, 0)
    _required_framework_version = (2, 0, 0)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        for plugin_dir in constants.PLUGINS_PATH:
            sids_json_file_name = os.path.join(
                plugin_dir, os.path.join("windows", "sids_and_privileges.json")
            )
            if os.path.exists(sids_json_file_name):
                break
        else:
            vollog.log(
                constants.LOGLEVEL_VVV,
                "sids_and_privileges.json file is missing plugin error",
            )
            raise RuntimeError(
                "The sids_and_privileges.json file missed from you plugin directory"
            )

        # Get all the sids from the json file.
        with open(sids_json_file_name, "r") as file_handle:
            sids_json_data = json.load(file_handle)
            self.servicesids = sids_json_data["service sids"]
            self.well_known_sids = sids_json_data["well known"]

            # Compile all the sids regex.
            self.well_known_sid_re = [
                (re.compile(c_list[0]), c_list[1])
                for c_list in sids_json_data["sids re"]
            ]

    @classmethod
    def get_requirements(cls) -> List[interfaces.configuration.RequirementInterface]:
        return [
            requirements.ModuleRequirement(
                name="kernel",
                description="Windows kernel",
                architectures=["Intel32", "Intel64"],
            ),
            requirements.ListRequirement(
                name="pid",
                description="Filter on specific process IDs",
                element_type=int,
                optional=True,
            ),
            requirements.PluginRequirement(
                name="pslist", plugin=pslist.PsList, version=(2, 0, 0)
            ),
            requirements.PluginRequirement(
                name="hivelist", plugin=hivelist.HiveList, version=(1, 0, 0)
            ),
        ]

    def lookup_user_sids(self) -> Dict[str, str]:
        """
        Enumerate the registry for all the users.

        Returns:
            An dictionary of {sid: user name}
        """

        key = "Microsoft\\Windows NT\\CurrentVersion\\ProfileList"
        val = "ProfileImagePath"
        kernel = self.context.modules[self.config["kernel"]]

        sids = {}
        for hive in hivelist.HiveList.list_hives(
            context=self.context,
            base_config_path=self.config_path,
            layer_name=kernel.layer_name,
            symbol_table=kernel.symbol_table_name,
            filter_string="config\\software",
            hive_offsets=None,
        ):
            try:
                for subkey in hive.get_key(key).get_subkeys():
                    sid = str(subkey.get_name())
                    path = ""
                    for node in subkey.get_values():
                        try:
                            value_node_name = node.get_name() or "(Default)"
                        except (
                            exceptions.InvalidAddressException,
                            layers.registry.RegistryFormatException,
                        ) as excp:
                            continue
                        try:
                            value_data = node.decode_data()
                            if isinstance(value_data, int):
                                value_data = format_hints.MultiTypeData(
                                    value_data, encoding="utf-8"
                                )
                            elif (
                                registry.RegValueTypes(node.Type)
                                == registry.RegValueTypes.REG_BINARY
                            ):
                                value_data = format_hints.MultiTypeData(
                                    value_data, show_hex=True
                                )
                            elif (
                                registry.RegValueTypes(node.Type)
                                == registry.RegValueTypes.REG_MULTI_SZ
                            ):
                                value_data = format_hints.MultiTypeData(
                                    value_data, encoding="utf-16-le", split_nulls=True
                                )
                            else:
                                value_data = format_hints.MultiTypeData(
                                    value_data, encoding="utf-16-le"
                                )
                            if value_node_name == val:
                                path = str(value_data).replace("\\x00", "")[:-1]
                                user = ntpath.basename(path)
                                sids[sid] = user
                        except (
                            ValueError,
                            exceptions.InvalidAddressException,
                            layers.registry.RegistryFormatException,
                        ) as excp:
                            continue
            except (KeyError, exceptions.InvalidAddressException):
                continue

        return sids

    def _generator(self, procs):
        user_sids = self.lookup_user_sids()

        # Go all over the process list, get the token
        for task in procs:
            # Make sure we have a valid token
            try:
                token = task.Token.dereference().cast("_TOKEN")
            except exceptions.InvalidAddressException:
                token = False

            if not token or not isinstance(token, interfaces.objects.ObjectInterface):
                yield (
                    0,
                    [
                        int(task.UniqueProcessId),
                        str(task.ImageFileName),
                        "Token unreadable",
                        "",
                    ],
                )
                continue

            # Go all over the sids and try to translate them with one of the tables we have
            for sid_string in token.get_sids():
                if sid_string in self.well_known_sids:
                    sid_name = self.well_known_sids[sid_string]
                elif sid_string in self.servicesids:
                    sid_name = self.servicesids[sid_string]
                elif sid_string in user_sids:
                    sid_name = user_sids[sid_string]
                else:
                    sid_name_re = find_sid_re(sid_string, self.well_known_sid_re)
                    if sid_name_re:
                        sid_name = sid_name_re
                    else:
                        sid_name = ""

                yield (
                    0,
                    (
                        task.UniqueProcessId,
                        objects.utility.array_to_string(task.ImageFileName),
                        sid_string,
                        sid_name,
                    ),
                )

    def run(self):
        filter_func = pslist.PsList.create_pid_filter(self.config.get("pid", None))
        kernel = self.context.modules[self.config["kernel"]]

        return renderers.TreeGrid(
            [("PID", int), ("Process", str), ("SID", str), ("Name", str)],
            self._generator(
                pslist.PsList.list_processes(
                    context=self.context,
                    layer_name=kernel.layer_name,
                    symbol_table=kernel.symbol_table_name,
                    filter_func=filter_func,
                )
            ),
        )


import bson
import ujson
from bson import ObjectId

from .utils import db

def retrieve_info(object_id):
    # Retrieve info from DB
    mongo = db.MongoDBConnection()
    result = list()

    with mongo:
        database = mongo.connection['mydb']
        collection = database['registrations']

        if object_id is not None:
            # Retrieve particular object
            try:
                single_object = collection.find_one({"_id": ObjectId(object_id)})
                if single_object is None:
                    return {"error": "Object id does not exist!"}
            except bson.errors.InvalidId:
                # catch the exception calling ObjectId
                return {"err": "Invalid object_id"}

            return {
                'id': object_id,
                'first_name': single_object['first_name'],
                'last_name': single_object['last_name'],
                'email': single_object['email']
            }

        else:
            # Retrieve all info
            for user in collection.find():
                result.append({
                    'id': str(user["_id"]),
                    'first_name': user['first_name'],
                    'last_name': user['last_name'],
                    'email': user['email']
                })
        return result

def lambda_handler(event, context):
    try:
        print(event)
        object_id = event['pathParameters']['id']
        # print("#######OBJECT_ID#########\n" + object_id)
    except TypeError:
        object_id = None
    except KeyError:
        object_id = None

    try:
        return {
            "statusCode": 200,
            "body": ujson.dumps({
                "message": "Success",
                "data": retrieve_info(object_id)
            })
        }
    except Exception as err:
        return {
            "statusCode": 400,
            "body": ujson.dumps({
                "message": "Something went wrong. Unable to parse data !",
                "error": str(err)
            })

        }


from flask import Flask, render_template, request, redirect, flash, session, config
from flask_debugtoolbar import DebugToolbarExtension
from surveys import satisfaction_survey as survey

app = Flask(__name__)
app.config['SECRET_KEY'] = "austin"
debug = DebugToolbarExtension(app)

responses = []
currentNum = 0
RESPONSES_KEY = "responses"

@app.route('/')
def show_home():
    """Main home page viewing."""
    global responses
    answers = responses
    responses = []
    instruction_text = str(survey.instructions)
    return render_template('base.html', instruction_text=instruction_text, answers=answers)

@app.route('/questions/1')
def question1():
    if len(responses) == 0:
        answers = responses
        global currentNum
        currentNum = 1
        question = survey.questions[0].question
        choices = survey.questions[0].choices
        return render_template('questions.html', question=question, num=currentNum, choices=choices, answers=answers)
    else:
        flash(f"Invalid question ID. Please take the survey in sequential order.")
        return redirect("/")

@app.route('/questions/2')
def question2():
    if len(responses) == 1:
        answers = responses
        global currentNum
        currentNum = 2
        question = survey.questions[1].question
        choices = survey.questions[1].choices
        return render_template('questions.html', question=question, num=currentNum, choices=choices, answers=answers)
    else:
        flash(f"Invalid question ID. Please take the survey in sequential order.")
        return redirect("/")

@app.route('/questions/3')
def question3():
    if len(responses) == 2:
        answers = responses
        global currentNum
        currentNum = 3
        question = survey.questions[2].question
        choices = survey.questions[2].choices
        return render_template('questions.html', question=question, num=currentNum, choices=choices, answers=answers)
    else:
        flash(f"Invalid question ID. Please take the survey in sequential order.")
        return redirect("/")

@app.route('/questions/4')
def question4():
    if len(responses) == 3:
        answers = responses
        global currentNum
        currentNum = 4
        question = survey.questions[3].question
        choices = survey.questions[3].choices
        return render_template('questions.html', question=question, num=currentNum, choices=choices, answers=answers)
    else:
        flash(f"Invalid question ID. Please take the survey in sequential order.")
        return redirect("/")

@app.route('/questions/5')
def question5():
    return redirect('/thankyou')

@app.route('/thankyou')
def thankyou_page():
    answers = responses
    session[RESPONSES_KEY] = responses
    return render_template('thankyou.html', answers=answers)

@app.route("/answer", methods=["POST"])
def post_answer():
    choice = request.form['answer']
    responses.append(choice)
    global currentNum
    newPageNum = currentNum + 1
    # return redirect('/questions/{newPageNum}')
    return redirect(f"/questions/{newPageNum}")


import os
# force TF to use CPU instead of GPU (sadly my discrete card is not support the latest CUDA version)
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

import gym

from gym import envs
import numpy as np
from keras.models import Sequential
from keras.layers import Dense
from keras.optimizers import Adam
import matplotlib.pyplot as plt

import time

# build DQN model from keras layers
def dqn_model(input_layer_size, output_layer_size, learning_rate=0.001):
    model = Sequential()
    model.add(Dense(128, input_dim=input_layer_size, activation='relu'))
    model.add(Dense(64, activation='relu'))
    model.add(Dense(output_layer_size, activation='linear'))
    model.compile(loss='mse', optimizer=Adam(lr=learning_rate))

    return model

if __name__ == "__main__":
    print(envs.registry.all())

    env = gym.make('LunarLander-v2')

    print(env.action_space)
    print(env.observation_space)

    OBSERVATION_SPACE_SIZE = env.observation_space.shape[0]
    ACTION_SPACE_SIZE = env.action_space.n

    model = dqn_model(OBSERVATION_SPACE_SIZE, ACTION_SPACE_SIZE)

    NUMBER_OF_EPISODES = 1000
    MAX_NUMBER_OF_STEPS_IN_EPISODE = 750

    BATCH_SIZE = 64

    DISCOUNT_RATE = 0.99

    exploration_rate = 1.0
    EXPLORATION_RATE_MINIMUM = 0.1
    EXPLORATION_RATE_CHANGE = (exploration_rate - EXPLORATION_RATE_MINIMUM) / ((NUMBER_OF_EPISODES / 2) * 100)

    obs_arr = np.array([])
    next_obs_arr = np.array([])
    actions_arr = np.array([], dtype=int)
    rewards_arr = np.array([])
    done_arr = np.array([])

    total_rewards = np.array([])

    for episode_idx in range(NUMBER_OF_EPISODES):
        start_time = time.time()
        # Restart the environment to start a new episode
        obs = env.reset()
        obs = np.reshape(obs, (1, OBSERVATION_SPACE_SIZE))

        episode_total_reward = 0

        explore_count = 0
        exploit_count = 0

        for step_idx in range(MAX_NUMBER_OF_STEPS_IN_EPISODE):
            # render a view with the environment current state
            # env.render()

            # choose action greedily (exploration-exploitation trade-off)
            if np.random.random() < exploration_rate:
                # explore
                explore_count += 1
                action = env.action_space.sample()
            else:
                # exploit: ask DQN to make a prediction of which action is preferred here
                exploit_count += 1
                prediction = model.predict(obs)[0]
                # print("prediction: ", prediction)
                action = np.argmax(prediction)

            # do the action
            next_obs, reward, is_done, info = env.step(action)
            next_obs = np.reshape(next_obs, (1, OBSERVATION_SPACE_SIZE))

            episode_total_reward += reward

            if len(obs_arr) == 0:
                obs_arr = np.array(obs)
                next_obs_arr = np.array(next_obs)
            else:
                obs_arr = np.vstack((obs_arr, obs))
                next_obs_arr = np.vstack((next_obs_arr, next_obs))
            actions_arr = np.append(actions_arr, action)
            rewards_arr = np.append(rewards_arr, reward)
            done_arr = np.append(done_arr, is_done)

            obs = next_obs

            if len(obs_arr) >= BATCH_SIZE:
                # print(episode_idx, ":", step_idx)
                indexes = np.random.randint(len(done_arr), size=BATCH_SIZE)

                batch_obs = np.squeeze(obs_arr[indexes])
                batch_next_obs = np.squeeze(next_obs_arr[indexes])
                batch_action = actions_arr[indexes]
                batch_reward = rewards_arr[indexes]
                batch_done = done_arr[indexes]

                all_targets = model.predict_on_batch(batch_obs)
                targets = batch_reward + DISCOUNT_RATE * (np.amax(model.predict_on_batch(batch_next_obs), axis=1)) * (1 - batch_done)

                all_targets[[np.array([i for i in range(BATCH_SIZE)])], batch_action] = targets

                # train the model
                model.fit(batch_obs, all_targets, verbose=0)

            # change exploration rate at every step until it becomes less or eq than exploration_rate_min
            # if exploration_rate > EXPLORATION_RATE_MINIMUM:
            #     exploration_rate -= EXPLORATION_RATE_CHANGE

            if exploration_rate > 0.01:
                exploration_rate *= 0.99

            # print("ep#%d:step#%d, obs=%s, reward=%s, done=%s" % (episode_idx, step_idx, observation, reward, is_done))
            if is_done:
                break

        total_rewards = np.append(total_rewards, episode_total_reward)

        print("----------------------")
        print("Episode %d reward: %d" % (episode_idx, episode_total_reward))
        print("Last step: %d" % (step_idx))
        print("Exploration rate: %f" % (exploration_rate))
        print("Explore count: %d" % (explore_count))
        print("Exploit count: %d" % (exploit_count))
        print("Time for episode: ", (time.time() - start_time), " sec.")

        # lunar lander problem is considered solved when total reward is 200+ points;
        # so we check if the last episodes was successful, the training would be stopped
        if np.mean(total_rewards[-100:]) >= 200:
            break

    # plot training progress
    x = list(range(len(total_rewards)))
    plt.plot(x, total_rewards)
    plt.xlabel("Episode")
    plt.ylabel("Total reward")
    plt.show()

    # save the model for future testing
    model.save('model.h5')

    # close gym environment
    env.close()