"""
Che-Hao Kang
University Bonn

Extract domain names from Malware
"""

#encoding=utf-8
from __future__ import unicode_literals
from ctypes import *
import ctypes
import collections
import re
import os
from winappdbg import win32
import string
import win32con
import win32api
import winerror
import sys
import gc
from ctypes import *
from time import time, sleep
from ctypes.wintypes import *
from threading import Thread
from struct import calcsize, pack, unpack, Struct
from win32gui import PyGetString, PySetString
from win32process import GetProcessMemoryInfo
from win32api import GetCurrentProcessId, GetCurrentProcess, CloseHandle
from win32con import MEM_FREE, PROCESS_VM_READ, PROCESS_VM_WRITE, PROCESS_QUERY_INFORMATION
from win32con import PAGE_READWRITE, PAGE_WRITECOPY, PAGE_EXECUTE_READWRITE, PAGE_EXECUTE_WRITECOPY
from win32con import PAGE_EXECUTE_READ, PAGE_READONLY, PROCESS_VM_OPERATION, PROCESS_ALL_ACCESS
import os
import shutil
import zipfile
import urllib
from lxml import html
import requests

# Get current processes in the operating system
def process_list():
    hProcessSnap = win32.CreateToolhelp32Snapshot()
    pe32 = win32.Process32First(hProcessSnap)
    if pe32 == None:
        print >> sys.stderr, "Failed getting first process."
        return

    while True:
        # szExeFile
        #    The name of the executable file for the process
        # https://www.ibm.com/developerworks/cn/opensource/os-cn-python-yield/
        hProcess = -1
        try:
            hProcess = win32.OpenProcess(win32.PROCESS_QUERY_INFORMATION | win32.PROCESS_VM_READ,
                                         False, pe32.th32ProcessID)
        except:
            pass
            #print "Get process handle FAILED:", ctypes.WinError()

        # return process, processID and size
        yield hProcess, pe32.th32ProcessID, pe32.szExeFile  # print pe32.szExeFile
        pe32 = win32.Process32Next(hProcessSnap)
        if pe32 == None:
            break

# Windows' memory-related things
MEMORY_STATES = {0x1000: "MEM_COMMIT", 0x10000: "MEM_FREE", 0x2000: "MEM_RESERVE"}
MEMORY_PROTECTIONS = {0x10: "PAGE_EXECUTE", 0x20: "PAGE_EXECUTE_READ", 0x40: "PAGE_EXECUTE_READWRITE",
                      0x80: "PAGE_EXECUTE_WRITECOPY", 0x01: "PAGE_NOACCESS", 0x02: "PAGE_READONLY",
                      0x04: "PAGE_READWRITE", 0x08: "PAGE_WRITECOPY"}
MEMORY_TYPES = {0x1000000: "MEM_IMAGE", 0x40000: "MEM_MAPPED", 0x20000: "MEM_PRIVATE"}

# Memory information structure
class PyMEMORY_BASIC_INFORMATION:
    def __init__(self, MBI):
        self.MBI = MBI
        self.set_attributes()

    def set_attributes(self):
        self.BaseAddress = self.MBI.BaseAddress
        self.AllocationBase = self.MBI.AllocationBase
        self.AllocationProtect = MEMORY_PROTECTIONS.get(self.MBI.AllocationProtect, self.MBI.AllocationProtect)
        self.RegionSize = self.MBI.RegionSize
        self.State = MEMORY_STATES.get(self.MBI.State, self.MBI.State)
        # self.Protect = self.MBI.Protect  # uncomment this and comment next line if you want to do a bitwise check on Protect.
        self.Protect = MEMORY_PROTECTIONS.get(self.MBI.Protect, self.MBI.Protect)
        self.Type = MEMORY_TYPES.get(self.MBI.Type, self.MBI.Type)

# Map memory information into PyMEMORY_BASIC_INFORMATION
def VirtualQueryEx(hProcess, lpAddress):
    _VirtualQueryEx = windll.kernel32.VirtualQueryEx
    _VirtualQueryEx.argtypes = [HANDLE, LPVOID, win32.PMEMORY_BASIC_INFORMATION, win32.SIZE_T]
    _VirtualQueryEx.restype  = win32.SIZE_T

    lpBuffer  = win32.MEMORY_BASIC_INFORMATION()
    dwLength  = sizeof(win32.MEMORY_BASIC_INFORMATION)
    success   = _VirtualQueryEx(hProcess, lpAddress, byref(lpBuffer), dwLength)
    if success == 0:
        raise ctypes.WinError()
    return PyMEMORY_BASIC_INFORMATION(lpBuffer)

# Scan a page of a process
def scan_page(process_handle, page_address):
    try:
        information = VirtualQueryEx(process_handle, page_address)
    except:
        #print "VirtualQueryEx FAILED:", ctypes.WinError()
        return "", False, "", "", ""

    base_address = information.BaseAddress
    access_protection = information.Protect
    region_size = information.RegionSize
    next_region = base_address + region_size

    if access_protection not in ["PAGE_EXECUTE_READWRITE"]:
        return next_region, False, base_address, region_size, access_protection
    else:
        # Extract all kinds of permissions and return
        return next_region, True, base_address, region_size, access_protection

# 1. Scan a process and call scan_page to get information from a page
# 2. A process' information is saved in the dictionary called addrePermSizeDict (address, size and permission)
# 3. If found RWX permission, return FoundRWX=True for further use
def scan_process(process_handle):
    # scans a processes pages for the target value.
    si = win32.SYSTEM_INFO()
    psi = byref(si)
    windll.kernel32.GetSystemInfo(psi)
    # get the first address of the first page to scan so we know where to start scanning
    # A pointer to the lowest memory address accessible to applications and dynamic-link libraries (DLLs)
    base_address = si.lpMinimumApplicationAddress
    # get the last address to scan so we know when to stop scanning.
    # A pointer to the highest memory address accessible to applications and DLLs
    max_address = si.lpMaximumApplicationAddress
    addrePermSizeDict = {}
    page_address = base_address
    FoundRWX = False

    while page_address < max_address:
        next_page, found, address, regionSize, permission = scan_page(process_handle, page_address)

        if address == "":
            continue

        permissionSize = []
        permissionSize.append(permission)
        permissionSize.append(regionSize)
        addrePermSizeDict[address] = permissionSize

        if found == True:
            FoundRWX = True

        page_address = next_page

    gc.collect()
    return addrePermSizeDict, FoundRWX

# Use regular expression to search valid domain names
def RECheckingDomain(content):
    urlFile = open('validUrls.txt', 'a')
    domainListFile = open('domainList.txt', 'r') # this list has all TLDs
    domainList = []
    for domain in domainListFile:
        domainList.append(domain.rstrip())

    # ((http://|https://)?:\/\/(?:www\.|(?!www))[^\s\.]+\.[^\s]{2,}|www\.[^\s]+\.[^\s]{2,})
    reString = '(http://|https://)?([a-zA-Z0-9][a-zA-Z0-9\_\-]{1,62}\.)+([a-zA-Z\-]{2,})'# + ([^a-zA-Z0-9/\.]|\n|$)'
    # reString = '/^[a-zA-Z0-9][a-zA-Z0-9-]{1,61}[a-zA-Z0-9]\.[a-zA-Z]{2,}$/'
    reobj = re.compile(reString, re.IGNORECASE | re.MULTILINE)
    m = reobj.finditer(content)

    # Use TLD list to check if the extracted domain name is valid
    for i in m:
        TLD = (i.group(3)).upper()
        Valid = False
        for domain in domainList:
            if TLD == domain:
                Valid = True
                break

        if Valid == True:
            urlFile.write("Valid: " + i.group() + "\n")

    urlFile.close()


# To check if PE exists in this process.
# If yes, dump the whole memory of this process to see whether domain names exist or not.
def dumpMemory(processInfoList):
    for proInfo in processInfoList:
        # Check if this process has RWX region.
        # If yes, go forward to see if PE files exist
        if not proInfo.processRWXFound:
            continue

        foundPE = False
        PEaddress = []

        # Check PE file
        for (address, permSize) in proInfo.processAddresses.items():
            if permSize[0] not in ["PAGE_EXECUTE_READWRITE"]:
                continue

            data = win32.ReadProcessMemory(proInfo.processHandle, address, 2)
            if len(data) != 2:
                raise ctypes.WinError()
            else:
                if data[0] == 'M' and data[1] == 'Z':
                    if not foundPE:
                        foundPE = True
                        print "\nFound PE in:", proInfo.processName, proInfo.processID

                        PEaddress.append(address)

                        urlFile = open('validUrls.txt', 'a')
                        urlFile.write('\n')
                        urlFile.write("\n" + str(proInfo.processName) + " " + str(proInfo.processID) + "\n")
                        urlFile.flush()
                        urlFile.close()
                        break

                    PEaddress.append(address)

                    data = win32.ReadProcessMemory(proInfo.processHandle, address, permSize[1])
                    RECheckingDomain(data)

        # if PE file found, then go to other regions
        if PEaddress:
            print "Start to extract urls by going through the whole region:", proInfo.processName, proInfo.processID
            print "Traversing......"

            for (address, size) in proInfo.processAddresses.items():
                try:
                    data = win32.ReadProcessMemory(proInfo.processHandle, address, permSize[1])
                    RECheckingDomain(data)
                except:
                    pass
                    # print ctypes.WinError()
            print "___ End of traversing", proInfo.processName, proInfo.processID, "___"

# 1. Eliminate duplicate urls
# 2. Put the original url file and modified url file into RESULTS folder
def checkDuplicateValidUrls():
    if not os.path.isfile('validUrls.txt'):
        print "\nNo valid urls found..." \
              "Please wait for 20 seconds to run it again!"
        return

    urlFile = open('validUrls.txt', 'r+')
    urlList = []

    for line in urlFile.readlines():
        if line.strip():
            urlList.append(line.strip())

    if os.path.isfile('validUrlsNoRep.txt'):
        os.remove("validUrlsNoRep.txt")

    urlNoRepFile = open('validUrlsNoRep.txt', 'a')

    # Delete duplicate urls
    urlList = list(set(urlList))

    for url in urlList:
        if url[0:6] == "Valid:":
            urlNoRepFile.write(url[7:].rstrip() + '\n')

    urlNoRepFile.close()
    urlFile.close()

    if not os.path.isdir('RESULTS'):
        os.mkdir('RESULTS')

    # Move files to RESULTS folder
    counter = 0
    counterNoRep = 0
    for root, dirs, files in os.walk("RESULTS/"):
        for f in files:
            filePath = os.path.join(root, f)
            index = string.find(filePath, '_validUrls.txt')
            indexNoRep = string.find(filePath, '_validUrlsNoRep.txt')
            if index != -1:
                counter += 1
            if indexNoRep != -1:
                counterNoRep += 1

    counter += 1
    shutil.move("validUrls.txt", str(counter) + "_validUrls.txt")
    shutil.move(str(counter) + '_validUrls.txt', 'RESULTS')

    counterNoRep += 1
    shutil.move("validUrlsNoRep.txt", str(counterNoRep) + "_validUrlsNoRep.txt")
    shutil.move(str(counterNoRep) + '_validUrlsNoRep.txt', 'RESULTS')

# Traverse RESULTS folder and combine urls together into allValidUrls.txt
def combineValidUrls():
    if not os.path.isdir('RESULTS'):
        print "No results found..."

    print "\nGenerate all valid urls from RESULTS folder"

    if os.path.isfile('allValidUrls.txt'):
        os.remove("allValidUrls.txt")

    allValidUrlsFile = open('allValidUrls.txt', 'a+')
    allValidUrlsList = []
    for root, dirs, files in os.walk("RESULTS/"):
        for f in files:
            filePath = os.path.join(root, f)
            if string.rfind(filePath, '_validUrlsNoRep.txt') != -1:
                validUrlsNoRepFile = open(filePath, 'r+')

                for line in validUrlsNoRepFile:
                    if line.strip():
                        line = line.rstrip()
                        same = False

                        for url in allValidUrlsList:
                            if line == url:
                                same = True
                                break

                        if same == False:
                            allValidUrlsList.append(line)

                validUrlsNoRepFile.close()

    allValidUrlsList.sort()
    for url in allValidUrlsList:
        allValidUrlsFile.write(url + '\n')

    allValidUrlsFile.close()

# Process Information Structure
class processInfo:
    processHandle = ""
    processID = ""
    processName = ""
    processAddresses = None
    #processAddressPerm = None
    processRWXFound = False


if __name__ == '__main__':
    if os.path.isfile('validUrls.txt'):
        os.remove("validUrls.txt")

    if os.path.isfile('validUrlsNoRep.txt'):
        os.remove("validUrlsNoRep.txt")

    delete = raw_input('Do you want to remove previous results? (y/n) : ')
    if delete == 'y' or delete == 'Y':
        if os.path.isdir('RESULTS'):
            shutil.rmtree('RESULTS')

    ######+++ Extract process information +++###
    processInfoList = []

    print "### Start to scan all processes to find PE files ### "

    for p in process_list():
        hProcess = p[0]
        processID = p[1]
        processName = p[2]

        if hProcess == -1:
            continue

        addrePermSizeDict, FoundRWX = scan_process(hProcess)
        if addrePermSizeDict:
            # Sort the dictionary by key
            addrePermSizeDict = collections.OrderedDict(sorted(addrePermSizeDict.items()))
            proIn = processInfo()
            proIn.processHandle = hProcess
            proIn.processID = processID
            proIn.processName = processName
            proIn.processAddresses = {}
            proIn.processAddresses = addrePermSizeDict
            proIn.processRWXFound = FoundRWX
            processInfoList.append(proIn)


    # To check if PE exists in this process.
    # If yes, dump the whole memory of this process to see whether domain names exist or not.
    dumpMemory(processInfoList)

    print "___ End of scanning all processes to find PE files ___"

    checkDuplicateValidUrls()
    combineValidUrls()
    ######--- Extract process information ---###