detector

#!/usr/bin/python3
# -*- coding: UTF-8 -*-
from pbox import *
from tinyscript import *
from tinyscript.report import *


__script__      = "Detection tool"
__version__     = "1.3.2"
__doc__         = """
This tool aims to detect the packer used on an input executable, folder of executables or Dataset.
"""
__description__ = "Detect the packer used on an input executable/folder/dataset"
__examples__    = [
    "my-sample.exe -d die",
    "/path/to/samples",
    "my-dataset -b -d die bintropy pypackerdetect",
]


def _metrics(tn=0, fp=0, fn=0, tp=0):
    """ Compute some metrics related to false/true positives/negatives. """
    accuracy  = float(tp + tn) / (tp + tn + fp + fn) if tp + tn + fp + fn > 0 else -1
    precision = float(tp) / (tp + fp) if tp + fp > 0 else -1
    recall    = float(tp) / (tp + fn) if tp + fn > 0 else -1                                      # or also sensitivity
    f_measure = 2. * precision * recall / (precision + recall) if precision + recall > 0 else -1  # or F(1)-score
    return accuracy, precision, recall, f_measure


if __name__ == '__main__':
    parser.add_argument("executable", help="executable or folder containing executables or dataset")
    parser.add_argument("-b", "--binary", action="store_true", help="only consider whether it is packed or not")
    parser.add_argument("-d", "--detector", action="extend", nargs="*", type=Detector.get, default=[],
                        help="detector(s) to be used", note="Uses the superdetector by default")
    grp = parser.add_mutually_exclusive_group()
    grp.add_argument("-f", "--failures-only", action="store_true", help="output failures only")
    grp.add_argument("-m", "--metrics-only", action="store_true", help="output metrics only")
    parser.add_argument("-n", "--not-labelled", action="store_true", help="count not labelled samples too")
    parser.add_argument("-s", "--do-not-strip-version", action="store_true",
                        help="do not strip the version number of the packer while applying detection")
    parser.add_argument("-t", "--threshold", type=float, help="threshold for voting when using multiple detectors")
    parser.add_argument("-w", "--weak", action="store_true", help="also consider weak assumptions")
    initialize(noargs_action="usage", multi_level_debug=True)
    configure_logging(args.verbose, ("INFO", "DEBUG", "DETAIL"))
    kw = {'debug': args.verbose > 0, 'verbose': args.verbose > 1, 'multiclass': not args.binary,
          'silent': args.verbose == 0, 'threshold': args.threshold, 'weak': args.weak}
    stats, cmatrix, pbar, errored = {}, {'fn': 0, 'fp': 0, 'tn': 0, 'tp': 0}, None, False
    single, mc = len(args.detector) == 1, kw['multiclass']
    logger.name = "detector" if single else "superdetector"
    kw['vote'] = v = len(args.detector) == 0
    single = len(args.detector) == 1
    kw['select'] = lst = [d for d in (args.detector or Detector.registry) if d.check(**kw)]
    if len(lst) == 0:
        sys.exit(0)
    NL, NP = [-1, NOT_LABELLED][mc], [0, NOT_PACKED][mc]
    if not single:
        logger.info("%s: %s" % (["Detectors", "Superdetector"][v], ", ".join(i.cname for i in lst)))
    try:
        with progress_bar(transient=True) as p:
            task = p.add_task("", total=None)
            for out in Detector.detect(args.executable, **kw):
                # ([exe]cutable, [r]esulting label of the detector, actual [l]abel, [d]etails) 
                details = {}
                if kw['debug']:
                    exe, r, l, details = out
                else:
                    exe, r, l = out
                if not args.do_not_strip_version and not args.binary:
                    l = strip_version(l)
                p._tasks[task].total = exe.len
                decision = "not labelled" if r == NL else "not packed" if r == NP else ["packed", r][mc]
                if not single:
                    logger.debug("final decision: %s" % decision)
                if not single and not args.metrics_only:
                    mlen = 0
                    for name in details.keys():
                        mlen = max(mlen, len(name))
                    msg = "Decisions:"
                    for name, label in details.items():
                        msg += ("\n{: <%d}: {}" % mlen).format(name, label)
                    if msg != "Decisions:":
                        logger.warning(msg)
                if not args.metrics_only and (not args.failures_only or r != l):
                    log = [logger.warning if r == NL else logger.failure if r != l else \
                           logger.success, logger.info][l == NL]
                    log("{}: {}".format(exe, decision))
                stats.setdefault(decision, 0)
                stats[decision] += 1
                if l != NL and args.not_labelled or decision != "not labelled":
                    cmatrix[['tp', 'tn'][r == NP] if r == l else ['fp', 'fn'][r == NP]] += 1
                p.update(task, advance=1.)
    except Exception as e:
        (logger.exception if args.verbose > 0 else logger.error)(e)
        if args.metrics_only:
            print("E,E,E,E")
        sys.exit(1)
    ph, h, m = metric_headers("classification"), ["Accuracy", "Precision", "Recall", "F-Measure"], _metrics(**cmatrix)
    m = [ph[k](v) if v >= 0 else "-" for k, v in zip(h, m)]
    ns, nc = sum(stats.values()), sum(cmatrix.values())
    rpt = [Section("Detection results: ")]
    if ns > 0:
        logger.debug("Statistics: %s" % stats)
    if nc > 0:
        logger.debug("Confusion Matrix: %s" % cmatrix)
    if args.metrics_only:
        print(",".join(m))
    elif ns == 1:
        print(r)
    else:
        if ns > 0:
            data = []
            for l, c in sorted(stats.items(), key=lambda x: x[0].lower()):
                data.append([l, c, "%.2f%%" % (100 * c / ns)])
            if len(data) > 0:
                rpt.append(Table(data, column_headers=["Label", "Count", "Percentage"]))
        if nc > 0:
            rpt.append(Table([m], column_headers=h))
        if ns + nc > 0:
            render(*rpt)