Malwarebytes.py

# The MIT License (MIT)

# Copyright (c) 2013, Nicholas Juszczak

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


# Script for automating the running of MalwareBytes
# quick scan, full scan, and update

# Last Modified: July 31, 2013


import pywinauto
import time
import sys
import logging
from pywinauto import application
from winsound import Beep

# Default values
args = {'executable':'',
        'update':False,
        'quick':False,
        'full':False }
# Logging
log = logging.getLogger()
ch = logging.StreamHandler(sys.stdout)
ch.setFormatter(logging.Formatter('%(message)s'))
log.addHandler(ch)
log.setLevel(logging.INFO)

# cli parameter parsing
if len(sys.argv) > 1:
    for index,arg in enumerate(sys.argv):
        if arg == '-e' or arg == '--executable':
            args['executable'] = sys.argv[index + 1]
        if arg == '-u' or arg == '--update':
            args['update'] = True
        if arg == '-q' or arg == '--quick':
            args['quick'] = True
            args['full'] = False
        if arg == '-f' or arg == '--full':
            args['full'] = True
            args['quick'] = False
        if arg == '-l' or arg == '--logfile':
            hdlr = logging.FileHandler(sys.argv[index + 1])
            formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
            hdlr.setFormatter(formatter)
            log.addHandler(hdlr)
            log.setLevel(logging.INFO)
        if arg == '-h' or arg == '--help':
            print 'Options are:'
            print '\t-e or --executable <path> | Path to MWB executable'
            print '\t-u or --update | Update MWB'
            print '\t-q or --quick | Quick scan (default)'
            print '\t-f or --full | Full scan'
            print '\t-l or --logfile <path> | log file'
            print '\t-h or --help | This screen'
            sys.exit(0)

print 'Remember, -h or --help for help!\n'
beepFreq = 2500
beepDur = 1000

def beep():
    try:
        Beep(beepFreq, beepDur)
    except:
        log.error('Unable to beep... :(')

# Run a update
def runUpdate():
    log.info("Updating MWB...")
    try:
        app = application.Application().start_(args['executable'] + ' /update')
    except:
        log.error('Error starting program...')
        beep()
        sys.exit(-1)

    # Attempt to connect to MWB
    dlg = None
    attempts = 30
    while dlg == None and attempts > 0:
        try:
            dlg = app.connect_(title_re="Updating .*")
        except:
            attempts -= 1
            time.sleep(2)

    if dlg == None:
        log.error("Error connecting to MalwareBytes")
        sys.exit(-1)
    else:
        log.info('Successfully connected to MWB')

    # Loop until update is finished
    attempts = 120
    successful = False
    while attempts > 0 and (not successful):
        # keep checking if close dialog has not appeared
        if not app.dlg.Static3.Exists():
            attempts -= 1
            time.sleep(2)
        # Error out if 'error' is found in dialog window
        elif 'error' in app.dlg.Static3.GetProperties()['Texts'][0]:
            log.error('Error while updating MWB...')
            beep()
            break
        # finish if done
        elif 'You have the latest' in app.dlg.Static3.GetProperties()['Texts'][0]:
            log.info('Already Updated')
            successful = True
            app.dlg.Button.Click()
            break
        elif 'successfully updated' in app.dlg.Static3.GetProperties()['Texts'][0]:
            log.info('Update Successful')
            successful = True
            app.dlg.Button.Click()
            break
        # Updating, keep waiting
        else:
            attempts -= 1
            time.sleep(2)
    
    #if not successfully updated
    if not successful:
	log.error('Failed to finish update')
        beep()
        sys.exit(-1)
    

#If scan was chosen
def runScan():
    if args['quick']:
        log.info('Running Quick Scan...')
        app = application.Application().start_(args['executable'] + " /scan")
    elif args['full']:
        log.info('Running Full Scan...')
        app = application.Application().start_(args['executable'] + " /fullscan")

    # Wait for window to be in correct state for 1 minute
    app['Malwarebytes Anti-Malware'].Wait('exists enabled', 60, 2)

    if not app['Malwarebytes Anti-Malware'].Exists():
        log.error("Error: Quick Scan timed out starting...")
        beep()
        sys.exit(-1)

    log.info('MWB scan started...')

    # Wait for Finish button to exist, be enabled. Wait for 4 hours, check 10 sec intervals
    app['Malwarebytes Anti-Malware'].Static3.Wait('exists enabled', 60*60*4, 10)

    if not app['Malwarebytes Anti-Malware'].Static3.Exists():
        log.error('Scan timed out scanning...')
        beep()
        sys.exit(-1)

    log.info('Popup found')

    # Check dialog of popup window for successful completion
    if 'scan completed successfully' in app['Malwarebytes Anti-Malware'].Static3.GetProperties()['Texts'][0]:
        log.info('Scan completed successfully')
        log.info('Attempting to close MWB...')



        app['Malwarebytes Anti-Malware'].OKButton.Wait('exists enabled', 5, 1)

        if not app['Malwarebytes Anti-Malware'].OKButton.Exists():
            log.error('Unable to click OK button')

        # If errors were shown
        if app['Malwarebytes Anti-Malware']['Show Results'].Exists():
            log.info('Showing results of scan')

            # continually try to click Show results to close app
            attempts = 5
            while app['Malwarebytes Anti-Malware']['Show Results'].Exists() and attempts > 0:
                app['Malwarebytes Anti-Malware']['Show Results'].Click()
                attempts -= 1
                time.sleep(1)

            # continually try to click remove selected to close app
            attempts = 5
            while app['Malwarebytes Anti-Malware']['Remove Selected'].Exists() and attempts > 0:
                app['Malwarebytes Anti-Malware']['Remove Selected'].Click()
                attempts -= 1
                time.sleep(1)

            # continually try to click yes to restart computer
            attempts = 5
            while app['Malwarebytes Anti-Malware'].Button.Exists() and attempts > 0:
                app['Malwarebytes Anti-Malware'].Button.Click()
                attempts -= 1
                time.sleep(1)

        # If either button still needs to be clicked we have a problem
        elif app['Malwarebytes Anti-Malware'].Exit.Exists():
            log.info('Exiting...')
            attempts = 5
            while app['Malwarebytes Anti-Malware'].Exit.Exists() and attempts > 0:
                app['Malwarebytes Anti-Malware'].Exit.Click()
                attempts -= 1
                time.sleep(1)

            if attempts == 0:
                log.error('Unable to close Exit button')
                beep()
                sys.exit(-1)


        log.info('MWB windows successfully closed')

        #Try to kill the notepad log that opens at the end of the scan
        try:
            app.connect_(title_re='.* Notepad')
            app.kill_()
            log.info('Notepad successfully closed')
        except:
            log.info('No Notepad instances opened...')
    # Else we've had an error scanning...
    else:
        log.error('Error Scanning...')
        beep()
        sys.exit(-1)


# Main
def main():
    if args['update']:
        runUpdate()
    if args['quick'] or args['full']:
        runScan()

if __name__ == '__main__':
    main()
    sys.exit(0)