handler.py

#!/usr/bin/env python3
import json
import zipfile
# import tarfile
# import gzip
import time
import rarfile
import os
import urllib3
import urllib.request
import urllib.parse
from bs4 import BeautifulSoup
import random

import logging
import boto3

logger = logging.getLogger()
logger.setLevel(logging.INFO)

def emptyDir(dirPath):

    # Empty /tmp folder for first time use
    for fileName in os.listdir(dirPath):

        os.remove(dirPath + fileName)

def zipDir(fileNames, dirPath, filePath):

    # Write a new zipfile without password
    with zipfile.ZipFile(filePath, 'w') as zipFile:        

        for fileName in fileNames:
            
            zipFile.write(dirPath + "/" + fileName)   

        zipFile.close()

    print("Zip successful")

def deleteFileWithFilePath(filePath):

    print("Deleting", filePath, "...")

    # Delete password protected zipfile    
    os.remove(filePath)

def randomChoiceDict(srcDict):

    return random.choice(list(srcDict.items()))

def randomChoiceList(srcList):

    return random.choice(srcList)

def getfile(http):

    files_url = "http://www.tekdefense.com/downloads/malware-samples/"
    download_base_url = "http://www.tekdefense.com"

    r = http.request('GET', files_url)

    parsed_html = BeautifulSoup(r.data, 'html.parser')

    tempList = parsed_html.body.find_all('h3', attrs={'class': 'title'})

    linksDict = {}

    for link in tempList:

        linksDict.update({link.text: download_base_url + urllib.parse.quote(link.find('a').get("href"))})

    tempDict = {}

    for link in linksDict.keys():

        if ".exe.zip" not in link:

            tempDict.update({link: linksDict[link]})
    
    return randomChoiceDict(tempDict)

def zipRepack(zipPassword, dirPath, fileName, filePath):
    
    with zipfile.ZipFile(filePath) as zipFile:

        zipFile.extractall(pwd=bytes(zipPassword, 'utf-8'), path=dirPath)       

    deleteFileWithFilePath(filePath=filePath)

    # fileNames = os.listdir(dirPath)
    # print("Directory - " + str(fileNames))
    
    # # if fileName in fileNames:
    # #     fileNames.remove(fileName)        
    
    # zipDir(fileNames, dirPath, filePath)

# def tarRepack(dirPath, fileName):
        
#     tar = tarfile.open(dirPath + fileName)
#     names = tar.getnames()
#     temp_file_path = ''
#     if os.path.isdir(dirPath + fileName + "_files"):
#         print('file already exist')
#         temp_file_path = os.path.isdir(dirPath + fileName + "_files")
#     else:
#         temp_file_path = os.mkdir(dirPath + fileName + "_files")
#         print('Create a new filename')
#     # Because there are many files after decompression, a directory with the same name should be established in advance
#     for name in names:
#         tar.extract(name, dirPath + fileName + "_files/")
#     tar.close()
#     return temp_file_path
    
# def gzipRepack(fileName):

#     f_name = fileName.replace(".gz", "")

#     # Get the name of the file, remove
#     g_file = gzip.GzipFile(fileName)

#     # Create gzip object
#     f = open(f_name, "w+").write(g_file.read())

#     # After the gzip object is opened with read(), it is written into the file created by open().
#     g_file.close()

def rarRepack(zipPassword, dirPath, fileName, filePath):
    
    with rarfile.Rarfile(dirPath + fileName) as file:
        file.extractall(pwd=zipPassword)
        file.close()

    deleteFileWithFilePath(filePath=filePath)

    # fileNames = os.listdir(dirPath)
    # zipDir(fileNames, dirPath, filePath)

def repackWithoutPassword(fileExtension, zipPassword, dirPath, fileName, filePath):
    
    if "zip" in fileExtension:

        zipRepack(zipPassword, dirPath, fileName, filePath)

    elif "tar" in fileExtension:

        rarRepack(zipPassword, dirPath, fileName, filePath)

    elif "gz" in fileExtension:

        rarRepack(zipPassword, dirPath, fileName, filePath)

    elif "rar" in fileExtension:
        
        rarRepack(zipPassword, dirPath, fileName, filePath)

    else:

        raise Exception("Unrecognized file extension: The file is not a zip, tar, tar.gz, rar.")

def s3WriteToFile(regionName, s3BucketName, dirPath, fileName):

    s3Client = boto3.client('s3', region_name=regionName)

    s3PutObjectResponse = s3Client.put_object(
        Bucket=s3BucketName,
        Key=fileName,
        Body=open(dirPath + "/" + fileName, 'rb'),
        ServerSideEncryption='AES256',
        StorageClass='REDUCED_REDUNDANCY'
    )

    print("\ns3PutObjectResponse - " + str(s3PutObjectResponse))

    # time.sleep(3)

    waiter = s3Client.get_waiter('object_exists')

    waiter.wait(        
        Bucket=s3BucketName,        
        Key=fileName
    )

    print("\nS3 ObjectExists!")

    return s3PutObjectResponse

def antimalwaretest(regionName, s3IngestBucketName):

    userAgent = {'user-agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/15.2 Safari/605.1.15'}
    http = urllib3.PoolManager(5, headers=userAgent)

    file = getfile(http)

    fileName = file[0]
    fileUrl = file[1]

    zipPassword = 'infected'

    # Attempt to download the various eicar test files
    print("---Running Anti-Malware Test---")
    print("Downloading Malware -", fileUrl)

    # Building a new filePath
    dirPath = "/tmp/"
    filePath = dirPath + fileName

    # Empty /tmp folder for first time use
    emptyDir(dirPath)

    # Download file to custom filePath
    urllib.request.urlretrieve(fileUrl, filePath)    

    fileExtension = fileName.split(".")[-1].lower()

    repackWithoutPassword(fileExtension, zipPassword, dirPath, fileName, filePath)    

    for fileName in os.listdir(dirPath):

        s3WriteToFile(regionName, s3IngestBucketName, dirPath, fileName)

    print("Remove all Malware Files")
    
    emptyDir(dirPath)

def main(event, context):

    regionName = str(os.environ.get("awsRegion"))
    s3IngestBucketName = str(os.environ.get("s3IngestBucketName"))

    if s3IngestBucketName[-1] == ",":
        s3IngestBucketName = s3IngestBucketName[:-1].replace(" ", "").split(",")
    else:
        s3IngestBucketName = s3IngestBucketName.replace(" ", "").split(",")

    print("Target S3 Bucket:", str(randomChoiceList(s3IngestBucketName)))

    antimalwaretest(regionName, randomChoiceList(s3IngestBucketName))

    body = {
        "message": "Go FSS-MalwareTest! Your function executed successfully!",
        "input": event 
    }

    return {"statusCode": 200, "body": json.dumps(body)}