main.py

#!/usr/bin/env python
# encoding: utf-8
"""
@author: w-devin
@contact: Mr.wyw@foxmail.com
@file: main.py
@time: 10/24/2020 01:50
@desc:
"""

import os
import sys
import click

from constant import NULL_STRING

from malware_detector.detector import Detector
from malware_detector.constant import *
from malware_detector.log_printer import *

detector = Detector()

def res_handler(res):
    if not res:
        return False

    bingo(res)
    return True

@click.command()
@click.option('-f', '--file2scan', type=click.STRING, default=NULL_STRING)
@click.option('-d', '--dir2scan', type=click.STRING, default=NULL_STRING)
@click.option('-p', '--payload', type=click.STRING, default=NULL_STRING)
def run(file2scan, dir2scan, payload):
    if file2scan:
        if os.path.isfile(file2scan):
            res = detector.detect(file2scan)
            res_handler(res)
        else:
            error('input a file with -f')

    if dir2scan:
        if os.path.isdir(dir2scan):
            info('scan dir: {}'.format(dir2scan))

            for root, dirs, files in os.walk(dir2scan):
                for filename in files:
                    file2scan = os.path.join(root, filename)
                    if os.path.isfile(file2scan):
                        res = detector.detect(file2scan)
                        res_handler(res)
        else:
            error('input a dir with -d')

    if payload:
        res = detector.detect(payload, PAYLOAD)
        res_handler(res)

if __name__ == '__main__':
    if len(sys.argv) <= 1:
        error('--help to get help message')
        exit()

    run()