was.py

__description__ = 'WAS - Wait A Sec: tool for scanning every file contained in a USB drive for malware before opening. Trust nobody.'
__author__ = 'Fabio Baroni <fabio@pentest.guru> @Fabiothebest89'
__version__ = '1.0.0'
__date__ = '13/09/2016'

'''
HISTORY:

6/8/2016 start of the project, media detection function
8/8/2016 added file hashing function
10/8/2016 added Virus Total result retrieval function
12/8/2016 added Virus Total result parsing and console notification
14/8/2016 added functionality to save results as a CSV file
16/8/2016 added audio notification in Italian and English language
13/9/2016 version 1.0.0 publicly released

TODO:
- add support for more languages
- implement file-locking function
- implement file upload function for scanning files not already scanned by Virus Total
- create Windows binary for ease of use by Windows folks
- add Linux support
'''

import win32api
import win32con
import win32gui
import hashlib
import os
from ctypes import *
import requests
import csv
import json
import time
import datetime
import configparser # import ConfigParser if you are using Python 2
import pyglet

script_dir = os.path.dirname(__file__)
# read configuration file
config = configparser.ConfigParser()
config.read("was-config.ini")
api = config.get("VIRUS-TOTAL", "api-key")
lang = config.get("NOTIFICATIONS", "lang")
audio = config.get("NOTIFICATIONS", "audio")
lock = config.get("FILE-LOCKING", "lock")
# EN sounds
device_scanning = pyglet.media.load(script_dir + '/audio/device_scanning.mp3')
virus_en = pyglet.media.load(script_dir + '/audio/virus_en.mp3')
# IT sounds
dispositivo_scansione = pyglet.media.load(script_dir + '/audio/dispositivo_scansione.mp3')
virus_it = pyglet.media.load(script_dir + '/audio/virus_it.mp3')
#
# Device change events (WM_DEVICECHANGE wParam)
#
DBT_DEVICEARRIVAL = 0x8000
DBT_DEVICEQUERYREMOVE = 0x8001
DBT_DEVICEQUERYREMOVEFAILED = 0x8002
DBT_DEVICEMOVEPENDING = 0x8003
DBT_DEVICEREMOVECOMPLETE = 0x8004
DBT_DEVICETYPESSPECIFIC = 0x8005
DBT_CONFIGCHANGED = 0x0018

#
# type of device in DEV_BROADCAST_HDR
#
DBT_DEVTYP_OEM = 0x00000000
DBT_DEVTYP_DEVNODE = 0x00000001
DBT_DEVTYP_VOLUME = 0x00000002
DBT_DEVTYPE_PORT = 0x00000003
DBT_DEVTYPE_NET = 0x00000004

#
# media types in DBT_DEVTYP_VOLUME
#
DBTF_MEDIA = 0x0001
DBTF_NET = 0x0002

WORD = c_ushort
DWORD = c_ulong


class DEV_BROADCAST_HDR(Structure):
    _fields_ = [
        ("dbch_size", DWORD),
        ("dbch_devicetype", DWORD),
        ("dbch_reserved", DWORD)
    ]


class DEV_BROADCAST_VOLUME(Structure):
    _fields_ = [
        ("dbcv_size", DWORD),
        ("dbcv_devicetype", DWORD),
        ("dbcv_reserved", DWORD),
        ("dbcv_unitmask", DWORD),
        ("dbcv_flags", WORD)
    ]


def drive_from_mask(mask):
    n_drive = 0
    while 1:
        if (mask & (2 ** n_drive)):
            return n_drive
        else:
            n_drive += 1

def retrieve_vt_report(filename, md5, hashes):
    url = "https://www.virustotal.com/vtapi/v2/file/report"
    global api
    global lang
    global audio
    global lock
    count = 1
    to_scan = []
    i = datetime.datetime.now()
    file_to_open = "{day}-{month}-{year}_{hour}-{minute}_malware_scan.csv".format(day = i.day, month = i.month, year = i.year, hour = i.hour, minute = i.minute)
    headers = ["filename", "md5", "positives", "permalink"]
    with open(file_to_open, "a") as f:
        f_csv = csv.writer(f)
        f_csv.writerow(headers)
        for file_hash in md5:
            params = {"apikey": api, "resource": file_hash}
            r = requests.post(url, data=params)
            if r.status_code == requests.codes.ALL_OK:
                report = json.loads(r.text)
                print(report)
                print(report['response_code'])
                if report['response_code'] == 0:
                    to_scan.append(file_hash)
                else:
                    positives = report['positives']
                    if positives != 0:
                        if audio == "ON" and lang == "EN":
                            virus_en.play()
                        elif audio =="ON" and lang == "IT":
                            virus_it.play()
                        else:
                            pass
                        permalink = report['permalink']
                        for x, y in iter(hashes.items()):
                            if y == file_hash:
                                name = x
                        row = [name, file_hash, positives, permalink]
                        f_csv.writerow(row)
            if count % 4 == 0:  # sleeps in order to ensure a max of 4 requests/min in accordance to the public API limit
                time.sleep(60)
                count += 1
    os.startfile(file_to_open)

def md5_files(path, blocksize = 2**20):
    hashes = {}
    for root, dirs, files in os.walk(path):
        for file in files:
            file_path = os.path.join(root, file)
            print(file_path)
            with open(file_path, "rb") as f:
                data = f.read(blocksize)
                hasher = hashlib.md5(data) # it's important to create a new MD5 object for every file
                while data:
                    data = f.read(blocksize)
                    hasher.update(data)
                    hashes[file_path] = hasher.hexdigest()
    if hashes:
        retrieve_vt_report(list(hashes.keys()), list(hashes.values()), hashes)
    return hashes


class Notification:
    def __init__(self):
        message_map = {
            win32con.WM_DEVICECHANGE: self.onDeviceChange
        }

        wc = win32gui.WNDCLASS()
        hinst = wc.hInstance = win32api.GetModuleHandle(None)
        wc.lpszClassName = "DeviceChangeDemo"
        wc.style = win32con.CS_VREDRAW | win32con.CS_HREDRAW
        wc.hCursor = win32gui.LoadCursor(0, win32con.IDC_ARROW)
        wc.hbrBackground = win32con.COLOR_WINDOW
        wc.lpfnWndProc = message_map
        classAtom = win32gui.RegisterClass(wc)
        style = win32con.WS_OVERLAPPED | win32con.WS_SYSMENU
        self.hwnd = win32gui.CreateWindow(
            classAtom,
            "Device Change Demo",
            style,
            0, 0,
            win32con.CW_USEDEFAULT, win32con.CW_USEDEFAULT,
            0, 0,
            hinst, None
        )

    def onDeviceChange(self, hwnd, msg, wparam, lparam):
        #
        # WM_DEVICECHANGE:
        #  wParam - type of change: arrival, removal etc.
        #  lParam - what's changed?
        #    if it's a volume then...
        #  lParam - what's changed more exactly
        #
        dev_broadcast_hdr = DEV_BROADCAST_HDR.from_address(lparam)
        global audio
        global lang

        if wparam == DBT_DEVICEARRIVAL:
            if dev_broadcast_hdr.dbch_devicetype == DBT_DEVTYP_VOLUME:
                if audio == "ON" and lang == "EN":
                    device_scanning.play()
                elif audio == "ON" and lang == "IT":
                    dispositivo_scansione.play()
                else:
                    pass

                dev_broadcast_volume = DEV_BROADCAST_VOLUME.from_address(lparam)
                drive_letter_bit = drive_from_mask(dev_broadcast_volume.dbcv_unitmask)
                drive_letter = chr(ord("A") + drive_letter_bit)
                letter_addendum = ":\\"
                drive_letter_path = "".join([drive_letter, letter_addendum])
                print(drive_letter_path)
                md5_files(drive_letter_path)


        return 1

if __name__ == '__main__':
    w = Notification()
    win32gui.PumpMessages()