# -*- coding: utf-8 -*-
from __future__ import print_function
import os
import glob
import subprocess
import re
from jinja2 import Template
import lief
from io import StringIO
import sys

import struct
import marshal
import zlib
from uuid import uuid4 as uniquename
import requests
from flask import Flask, render_template, request
from flask_dropzone import Dropzone
from bs4 import BeautifulSoup
import jsbeautifier
from urllib.parse import urljoin
import checker

basedir = os.path.abspath(os.path.dirname(__file__))


app = Flask(__name__)

app.config.update(
    UPLOADED_PATH=os.path.join(basedir, 'uploads'),
    # Flask-Dropzone config:
    
    DROPZONE_MAX_FILE_SIZE=1024,
    DROPZONE_MAX_FILES=1,
    DROPZONE_REDIRECT_VIEW='completed'  # set redirect view
)

dropzone = Dropzone(app)

def url_finder(string): 
    regex = r"(?i)\b((?:https?://|www\d{0,3}[.]|[a-z0-9.\-]+[.][a-z]{2,4}/)(?:[^\s()<>]+|\(([^\s()<>]+|(\([^\s()<>]+\)))*\))+(?:\(([^\s()<>]+|(\([^\s()<>]+\)))*\)|[^\s`!()\[\]{};:'\".,<>?«»“”‘’]))"
    url = re.findall(regex,string)       
    return [x[0] for x in url] 

def ipaddr_finder(string):
  regex = r'(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})'
  ip = re.findall(regex,string)       
  return ip 


if sys.version_info.major == 3:
    from importlib.util import MAGIC_NUMBER
    pyc_magic = MAGIC_NUMBER
else:
    import imp
    pyc_magic = imp.get_magic()


class CTOCEntry:
    def __init__(self, position, cmprsdDataSize, uncmprsdDataSize, cmprsFlag, typeCmprsData, name):
        self.position = position
        self.cmprsdDataSize = cmprsdDataSize
        self.uncmprsdDataSize = uncmprsdDataSize
        self.cmprsFlag = cmprsFlag
        self.typeCmprsData = typeCmprsData
        self.name = name


class PyInstArchive:
    PYINST20_COOKIE_SIZE = 24           # For pyinstaller 2.0
    PYINST21_COOKIE_SIZE = 24 + 64      # For pyinstaller 2.1+
    MAGIC = b'MEI\014\013\012\013\016'  # Magic number which identifies pyinstaller

    def __init__(self, path):
        self.filePath = path


    def open(self):
        try:
            self.fPtr = open(self.filePath, 'rb')
            self.fileSize = os.stat(self.filePath).st_size
        except:
            print('[!] Error: Could not open {0}'.format(self.filePath))
            return False
        return True


    def close(self):
        try:
            self.fPtr.close()
        except:
            pass


    def checkFile(self):
        print('[+] Processing {0}'.format(self.filePath))
        # Check if it is a 2.0 archive
        self.fPtr.seek(self.fileSize - self.PYINST20_COOKIE_SIZE, os.SEEK_SET)
        magicFromFile = self.fPtr.read(len(self.MAGIC))

        if magicFromFile == self.MAGIC:
            self.pyinstVer = 20     # pyinstaller 2.0
            print('[+] Pyinstaller version: 2.0')
            return True

        # Check for pyinstaller 2.1+ before bailing out
        self.fPtr.seek(self.fileSize - self.PYINST21_COOKIE_SIZE, os.SEEK_SET)
        magicFromFile = self.fPtr.read(len(self.MAGIC))

        if magicFromFile == self.MAGIC:
            print('[+] Pyinstaller version: 2.1+')
            self.pyinstVer = 21     # pyinstaller 2.1+
            return True

        print('[!] Error : Unsupported pyinstaller version or not a pyinstaller archive')
        return False


    def getCArchiveInfo(self):
        try:
            if self.pyinstVer == 20:
                self.fPtr.seek(self.fileSize - self.PYINST20_COOKIE_SIZE, os.SEEK_SET)

                # Read CArchive cookie
                (magic, lengthofPackage, toc, tocLen, self.pyver) = \
                struct.unpack('!8siiii', self.fPtr.read(self.PYINST20_COOKIE_SIZE))

            elif self.pyinstVer == 21:
                self.fPtr.seek(self.fileSize - self.PYINST21_COOKIE_SIZE, os.SEEK_SET)

                # Read CArchive cookie
                (magic, lengthofPackage, toc, tocLen, self.pyver, pylibname) = \
                struct.unpack('!8siiii64s', self.fPtr.read(self.PYINST21_COOKIE_SIZE))

        except:
            print('[!] Error : The file is not a pyinstaller archive')
            return False

        print('[+] Python version: {0}'.format(self.pyver))

        # Overlay is the data appended at the end of the PE
        self.overlaySize = lengthofPackage
        self.overlayPos = self.fileSize - self.overlaySize
        self.tableOfContentsPos = self.overlayPos + toc
        self.tableOfContentsSize = tocLen

        print('[+] Length of package: {0} bytes'.format(self.overlaySize))
        return True


    def parseTOC(self):
        # Go to the table of contents
        self.fPtr.seek(self.tableOfContentsPos, os.SEEK_SET)

        self.tocList = []
        parsedLen = 0

        # Parse table of contents
        while parsedLen < self.tableOfContentsSize:
            (entrySize, ) = struct.unpack('!i', self.fPtr.read(4))
            nameLen = struct.calcsize('!iiiiBc')

            (entryPos, cmprsdDataSize, uncmprsdDataSize, cmprsFlag, typeCmprsData, name) = \
            struct.unpack( \
                '!iiiBc{0}s'.format(entrySize - nameLen), \
                self.fPtr.read(entrySize - 4))

            name = name.decode('utf-8').rstrip('\0')
            if len(name) == 0:
                name = str(uniquename())
                print('[!] Warning: Found an unamed file in CArchive. Using random name {0}'.format(name))

            self.tocList.append( \
                                CTOCEntry(                      \
                                    self.overlayPos + entryPos, \
                                    cmprsdDataSize,             \
                                    uncmprsdDataSize,           \
                                    cmprsFlag,                  \
                                    typeCmprsData,              \
                                    name                        \
                                ))

            parsedLen += entrySize
        print('[+] Found {0} files in CArchive'.format(len(self.tocList)))


    def _writeRawData(self, filepath, data):
        nm = filepath.replace('\\', os.path.sep).replace('/', os.path.sep).replace('..', '__')
        nmDir = os.path.dirname(nm)
        if nmDir != '' and not os.path.exists(nmDir): # Check if path exists, create if not
            os.makedirs(nmDir)

        with open(nm, 'wb') as f:
            f.write(data)


    def extractFiles(self):
        print('[+] Beginning extraction...please standby')
        extractionDir = os.path.join(os.getcwd(), os.path.basename(self.filePath) + '_extracted')

        if not os.path.exists(extractionDir):
            os.mkdir(extractionDir)

        os.chdir(extractionDir)

        for entry in self.tocList:
            basePath = os.path.dirname(entry.name)
            if basePath != '':
                # Check if path exists, create if not
                if not os.path.exists(basePath):
                    os.makedirs(basePath)

            self.fPtr.seek(entry.position, os.SEEK_SET)
            data = self.fPtr.read(entry.cmprsdDataSize)

            if entry.cmprsFlag == 1:
                data = zlib.decompress(data)
                # Malware may tamper with the uncompressed size
                # Comment out the assertion in such a case
                assert len(data) == entry.uncmprsdDataSize # Sanity Check

            if entry.typeCmprsData == b's':
                # s -> ARCHIVE_ITEM_PYSOURCE
                # Entry point are expected to be python scripts
                print('[+] Possible entry point: {0}.pyc'.format(entry.name))
                self._writePyc(entry.name + '.pyc', data)

            elif entry.typeCmprsData == b'M' or entry.typeCmprsData == b'm':
                # M -> ARCHIVE_ITEM_PYPACKAGE
                # m -> ARCHIVE_ITEM_PYMODULE
                # packages and modules are pyc files with their header's intact
                self._writeRawData(entry.name + '.pyc', data)

            else:
                self._writeRawData(entry.name, data)

                if entry.typeCmprsData == b'z' or entry.typeCmprsData == b'Z':
                    self._extractPyz(entry.name)


    def _writePyc(self, filename, data):
        with open(filename, 'wb') as pycFile:
            pycFile.write(pyc_magic)            # pyc magic

            if self.pyver >= 37:                # PEP 552 -- Deterministic pycs
                pycFile.write(b'\0' * 4)        # Bitfield
                pycFile.write(b'\0' * 8)        # (Timestamp + size) || hash 

            else:
                pycFile.write(b'\0' * 4)      # Timestamp
                if self.pyver >= 33:
                    pycFile.write(b'\0' * 4)  # Size parameter added in Python 3.3

            pycFile.write(data)


    def _extractPyz(self, name):
        dirName =  name + '_extracted'
        # Create a directory for the contents of the pyz
        if not os.path.exists(dirName):
            os.mkdir(dirName)

        with open(name, 'rb') as f:
            pyzMagic = f.read(4)
            assert pyzMagic == b'PYZ\0' # Sanity Check

            pycHeader = f.read(4) # Python magic value

            # Skip PYZ extraction if not running under the same python version
            if pyc_magic != pycHeader:
                print('[!] Warning: This script is running in a different Python version than the one used to build the executable.')
                print('[!] Please run this script in Python{0} to prevent extraction errors during unmarshalling'.format(self.pyver))
                print('[!] Skipping pyz extraction')
                return

            (tocPosition, ) = struct.unpack('!i', f.read(4))
            f.seek(tocPosition, os.SEEK_SET)

            try:
                toc = marshal.load(f)
            except:
                print('[!] Unmarshalling FAILED. Cannot extract {0}. Extracting remaining files.'.format(name))
                return

            print('[+] Found {0} files in PYZ archive'.format(len(toc)))

            # From pyinstaller 3.1+ toc is a list of tuples
            if type(toc) == list:
                toc = dict(toc)

            for key in toc.keys():
                (ispkg, pos, length) = toc[key]
                f.seek(pos, os.SEEK_SET)
                fileName = key

                try:
                    # for Python > 3.3 some keys are bytes object some are str object
                    fileName = fileName.decode('utf-8')
                except:
                    pass

                # Prevent writing outside dirName
                fileName = fileName.replace('..', '__').replace('.', os.path.sep)

                if ispkg == 1:
                    filePath = os.path.join(dirName, fileName, '__init__.pyc')

                else:
                    filePath = os.path.join(dirName, fileName + '.pyc')

                fileDir = os.path.dirname(filePath)
                if not os.path.exists(fileDir):
                    os.makedirs(fileDir)

                try:
                    data = f.read(length)
                    data = zlib.decompress(data)
                except:
                    print('[!] Error: Failed to decompress {0}, probably encrypted. Extracting as is.'.format(filePath))
                    open(filePath + '.encrypted', 'wb').write(data)
                else:
                    self._writePyc(filePath, data)


def pyintaller_module(filename):
    arch = PyInstArchive(filename)
    if arch.open():
        if arch.checkFile():
            if arch.getCArchiveInfo():
                arch.parseTOC()
                arch.extractFiles()
                arch.close()
                print('[+] Successfully extracted pyinstaller archive: {0}'.format(filename))
                print('')
                print('You can now use a python decompiler on the pyc files within the extracted directory')
                return
        arch.close()

@app.route('/url', methods=['GET', 'POST'])
def index():
    
    results = {}
    if request.method == "POST":
        # get url that the user has entered
        
        url = request.form['url']
        
        print(url)
        r = requests.get(url)
        soup = BeautifulSoup(r.content, features="html5lib")
        l = [i.get('src') for i in soup.find_all('script') if i.get('src')] 

        final_js_list = []

        for i in l:
            script_url = urljoin(url, i)
            final_js_list.append(script_url)

        for js_file in final_js_list:
            r = requests.get(js_file)

        res = jsbeautifier.beautify(
            r.content.decode('utf-8')
)

    return render_template('url_report.html', results=results)

@app.route('/', methods=['POST', 'GET'])
def upload():
    if request.method == 'POST':
        f = request.files.get('file')
        f.save(os.path.join(app.config['UPLOADED_PATH'], f.filename))
        
    return render_template('index.html')


def floss_module(file_path):
    
    latest_file = file_path

    FLOSS_PATH = "C:\\Users\\Kimaya\\Downloads\\floss-v1.7.0-windows\\floss.exe"
    try:
        output = subprocess.check_output(
                ["powershell.exe", FLOSS_PATH, 
                "--no-static-strings", latest_file], 
                shell=True)
    except subprocess.CalledProcessError as e:
        print ("subproces CalledProcessError.output =", e.output)
    # print (output.decode())

    logfile = open("floss_output.txt", 'w+')
    logfile.write(output.decode())
    logfile.close()

    
    decoded_strings = []
    skip_list = ['FLOSS decoded', 'FLOSS extracted', 'Finished execution']

    with open("floss_output.txt", 'r') as f:
        for line in f:
            line = line.rstrip('\n')
            if(len(line) != 0):
                if any([x in line for x in skip_list]):
                    pass
                else:
                    decoded_strings.append(line)
                    
    print(decoded_strings)

    temp_str1 = ""
    temp_str2 = ""

    for str in decoded_strings:
        urls_found = url_finder(str)
        if(len(urls_found) > 0):
            print("Obfuscated URL found!")
            str1 = temp_str1.join(urls_found)
                

        ips = ipaddr_finder(str)
        if(len(ips) > 0):
            print("Obfuscated ips found!")
            str2 = temp_str2.join(ips)

    ret_str = str1 +" "+ str2
    ret_str = ret_str.replace(" ",' <br> ')
    binary = lief.parse(latest_file)
    # print(dir(binary))

    
    old_stdout = sys.stdout
    new_stdout = StringIO()
    sys.stdout = new_stdout

    print(binary.header)

    output_header = new_stdout.getvalue()

    sys.stdout = old_stdout
    print(output_header)
    output_header = output_header.replace('\n',' <br> ')

    return ret_str, output_header


def uncompyle_runner(file_path, save_dir):
    UNC_PATH = "C:\\Users\\Kimaya\\Desktop\\BTech_Project\\uncompyle6.exe"
    try:
        output = subprocess.check_output(
                ["powershell.exe", UNC_PATH, 
                "-o", save_dir,file_path], 
                shell=True)
    except subprocess.CalledProcessError as e:
        print ("subproces CalledProcessError.output =", e.output)


def bandit_runner(directory):
    output = "ignore"
    try:
        output = subprocess.check_output(
                ["powershell.exe", "bandit", 
                "-r", directory], 
                shell=True)
    except subprocess.CalledProcessError as e:
        print ("subproces CalledProcessError.output =", e.output)
        
        print(output)
        logfile = open("bandit_output.txt", 'w+')
        logfile.write(e.output.decode())
        logfile.close()
        with open("bandit_output.txt") as f:
            t = f.read()

        ind1 = t.find("Test results:")
        ind2 = t.find("Code scanned:")
        index1 = t.find("Run metrics:")
        index2 = t.find("Files skipped")

        imp_part1 = t[ind1:ind2 - 3]

        list1 = list()
        list1 = imp_part1.split("\n")

        final_list = []
        for i in list1:
            if(i.startswith("   Location:") or i.startswith("   More Info") or (i == '')):
                continue
            else:
                final_list.append(i)
                

        final_list.append(t[index1:index2])
        # listToStr = ' '.join([str(elem) for elem in final_list])
        with open('bandit_final_output.txt', 'w') as filehandle:
	        filehandle.writelines("%s\n" % i for i in final_list)

        with open("bandit_final_output.txt") as f:
            txt1 = f.read()

        return txt1


                # with open('bandit_output.txt', 'w') as filehandle:
        #     data = filehandle.read().replace('\n', '')

    

def xor_module(file_path):
    output = "ignore"
    try:
        output = subprocess.check_output(
                ["powershell.exe", "xortool", file_path,
                "-c", "20"], 
                shell=True)
    except subprocess.CalledProcessError as e:
        print ("subproces CalledProcessError.output =", e.output)

    # print(output)

    logfile = open("xor_output.txt", 'w+')
    logfile.write(output.decode())
    logfile.close()

    return output



def py_extract_module(file_path):
    pyintaller_module(file_path)
    parent_dir, file_name  = os.path.split(file_path)
    grandparent_dir, xyz = os.path.split(parent_dir)
    print(grandparent_dir,xyz)
    os.mkdir(os.path.join(grandparent_dir,"{}_source".format(file_name)))

    list_of_files = glob.glob(os.path.join(grandparent_dir,"{}_extracted".format(file_name), '*.pyc'))
    # print(list_of_files)
    for file in list_of_files:
        uncompyle_runner(file,os.path.join(grandparent_dir,"{}_source".format(file_name)))

    

    return bandit_runner(os.path.join(grandparent_dir,"{}_source".format(file_name)))


  
@app.route('/completed')
def completed():
    list_of_files = glob.glob(os.path.join(app.config['UPLOADED_PATH'], '*'))
    latest_file = max(list_of_files, key=os.path.getctime)
    print(latest_file)
   

    if("tool_xored" in latest_file):
            
            xor_output = xor_module(latest_file)
            xor_output = xor_output.decode().replace('\n',' <br> ')

            # os.remove(latest_file)
            return render_template('xor_report.html', xor_str = xor_output)

    
    ml_output = checker.check(latest_file)

    if(ml_output == 'malicious'):
        ml_out_str = 'The given file is malicious'
        binary = lief.parse(latest_file)
        old_stdout = sys.stdout
        new_stdout = StringIO()
        sys.stdout = new_stdout
        print(binary.header)
        ml_info = new_stdout.getvalue()

        sys.stdout = old_stdout
        # print(outputh2)
        ml_info = ml_info.replace('\n',' <br> ')
        return render_template('ml_report.html', ml_Str = ml_out_str, info_Str = ml_info)

    else:



    # Strings
        if("cexe" in latest_file):
            ret_str, output_header = floss_module(latest_file)
            # os.remove(latest_file)
            return render_template('str_report.html', deobfs_Str = ret_str, info_Str = output_header)
        
        # XOR
        if("tool_xored" in latest_file):
            
            xor_output = xor_module(latest_file)
            xor_output = xor_output.decode().replace('\n',' <br> ')

            # os.remove(latest_file)
            return render_template('xor_report.html', xor_str = xor_output)

        # Python exe
        if("hello" in latest_file):
            py_output =  py_extract_module(latest_file)
            py_output = py_output.replace('\n',' <br> ')

            binary = lief.parse(latest_file)
            old_stdout = sys.stdout
            new_stdout = StringIO()
            sys.stdout = new_stdout
            print(binary.header)
            outputh2 = new_stdout.getvalue()

            sys.stdout = old_stdout
            # print(outputh2)
            outputh2 = outputh2.replace('\n',' <br> ')

            # os.remove(latest_file)
            return render_template('py_report.html', py_str = py_output,info_Str = outputh2)

    # else:
    #     py_output =  py_extract_module(latest_file)
    #     py_output = py_output.replace('\n',' <br> ')

    #     xor_output = xor_module(latest_file)
    #     xor_output = xor_output.decode().replace('\n',' <br> ')

    #     ret_str, output_header = floss_module(latest_file)

    #     return render_template('else_report.html', py_str = py_output,  xor_str = xor_output, deobfs_Str = ret_str)





    

if __name__ == '__main__':
    app.run(debug=True)
