hello-detector.py

#!/usr/bin/env python
"""
Hello Detector

Joe Vasquez <[email protected]>

For CYB5280, check supplied COM files for polymorphic
Hello World virus.

"""

from binascii import hexlify
from os import listdir, path


def get_COM_files(directory):
    """
    Get COM Files

    Look for COM files in direcotry, return list of paths.
    """
    files = []
    for element in listdir(directory):
        candidate = path.join(directory, element)
        if path.isfile(candidate):
            name, extension = path.splitext(candidate)
            if extension == '.COM' or extension == '.com':
                files.append(candidate)
    return files

def check_1(bytes):
    """
    Check 1

    Find the hex values in findItems list in order.
    Ignores junk this way.
    """
    findItems = ('09', 'cd', '21', 'b8', '4c', 'cd', '21')

    index = 0
    for findItem in findItems:
        position = bytes.find(findItem, index)
        if position == -1:
            return False
        index = position + len(findItem)

    return True

def check_2(bytes):
    """Find 'Hello world'"""
    bytes = bytes.replace('24', '')
    return 'Hello World!'.encode('hex') in bytes


if __name__ == "__main__":
    print u'running program\n'

    currentDirectory = path.dirname(path.realpath(__file__))

    print u'current directory:  %s' % currentDirectory

    comFiles = get_COM_files(currentDirectory)

    print u'found %s COM files' % len(comFiles)

    for comFile in comFiles:
        print u'-----------------------------'
        print u'checking file:  %s' % comFile
        size  = path.getsize(comFile)
        print u'size:  %s bytes' % size
        bytes = ''
        bytes_ascii = ''
        with open(comFile, 'rb') as f:
            bytes = f.read(size)

        bytes_ascii = hexlify(bytes)

        check1 = check_1(bytes_ascii)
        if not check1:
            print u'Check one failed, this is not the virus.'
        else:
            check2 = check_2(bytes_ascii)
            if not check2:
                print u'Check two failed, this is not the virus.'
            else:
                print u'This file is the virus!'