avred/avredhelper.py at main · dobin/avred · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
#!/usr/bin/python3

import os
import pickle
import argparse
import pstats
import r2pipe

from scanner import *
from model.model_base import *
from model.model_code import *
from plugins.pe.file_pe import FilePe
from plugins.pe.augment_pe import DataReferor
from myutils import getOutcomesFromDir, OutcomesToCsv

HASHCACHE_FILE = "hashcache.pickle"


def hashcache():
    if not os.path.exists(HASHCACHE_FILE):
        print("HashCache file does not exist: {}".format(HASHCACHE_FILE))

    with open(HASHCACHE_FILE, "rb") as file:
        cache = pickle.load(file)
        print("TimeRounded;Time;Filename;Scanner;Result")
        for entry in cache.values():
            if entry.scanTime > 1:
                scantime = round(entry.scanTime, 1)
                scantime = str(scantime).replace('.', ',')
                print("{};{};{};{};{}".format(scantime,entry.scanTime, entry.scannerName, entry.filename, entry.result))


def printoutcome(filename: str):
    with open(filename, "rb") as input_file:
        outcome = pickle.load(input_file)
        print(str(outcome))


def patchfile(fname: str, pos: int, data: bytes):
    print( f"Writing {len} bytes to file {fname} at position {pos} ")

    fp = open(fname, "r+b")

    fp.seek(pos)
    fp.write(data)

    fp.close()


def printFileInfo(filepath):
    filePe = FilePe()
    filePe.loadFromFile(filepath)

    print("Sections:")
    for section in filePe.peSectionsBag.sections:
        print(section)

    print("")
    print("Regions:")
    for region in filePe.regionsBag.sections:
        print(region)


def printFileDataInfo(filepath, offset, size):
    filePe = FilePe()
    filePe.loadFromFile(filepath)

    r2 = r2pipe.open(filePe.filepath)
    r2.cmd("aaa")  # aaaa
    dataReferor = DataReferor(r2)
    dataReferor.init()
    #for s in dataReferor.stringsIt:
    #    print(s[2])

    disasmLines = dataReferor.query(offset, size)
    for disasmLine in disasmLines:
        print(disasmLine)


def printperf():
    # python3 -m cProfile -s time -o perf.txt nkeyrollover.py
    p = pstats.Stats('perf.txt')
    p.sort_stats('cumulative').print_stats()


def printcsv(dir: str):
    outcomes = getOutcomesFromDir(dir)
    print(OutcomesToCsv(outcomes))


def main():
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", help="Choose a command")

    # Hash Cache
    parserHashcache = subparsers.add_parser("hashcache", help="")

    # CSV
    parserCsv = subparsers.add_parser("csv", help="")
    parserCsv.add_argument("--directory", help="")

    # Outcome
    parserOutcome = subparsers.add_parser("outcome", help="")
    parserOutcome.add_argument("--file", help="")

    # Info
    parserInfo = subparsers.add_parser("info", help="")
    parserInfo.add_argument("--file", help="")

    # Augment
    parserAugment = subparsers.add_parser("augment", help="")
    parserAugment.add_argument("--file", help="")
    parserAugment.add_argument("--offset", help="", type=int)
    parserAugment.add_argument("--size", help="", type=int)

    args = parser.parse_args()
    if args.command == "csv":
        printcsv(args.directory)
    elif args.command == "hashcache":
        hashcache()
    elif args.command == "info":
        printFileInfo(args.file)
    elif args.command == "augment":
        printFileDataInfo(args.file, args.offset, args.size)
    elif args.command == "outcome":
        printoutcome(args.file)


if __name__ == "__main__":
    main()