80 lines
3.0 KiB
Python
80 lines
3.0 KiB
Python
# !/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
|
|
import binascii
|
|
import os
|
|
from pybloom import BloomFilter
|
|
|
|
|
|
import argparse
|
|
|
|
|
|
|
|
# reference - http://stackoverflow.com/a/9631635
|
|
def blocks(this_file, size=65536):
|
|
while True:
|
|
b = this_file.read(size)
|
|
if not b:
|
|
break
|
|
yield b
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(prog='build.py')
|
|
parser.add_argument("-v", "--verbose", help="Display verbose output message", action="store_true", required=False)
|
|
config = parser.add_mutually_exclusive_group()
|
|
config.add_argument('-f', "--config", default="/nsrl/nsrl.conf",help='Config file with all settings')
|
|
settings = config.add_argument_group()
|
|
settings.add_argument('-e','--error-rate', type=float, default=0.01 ,help="Error Rate for False-Positives")
|
|
settings.add_argument('-n','--hashcount',type=int, help="Provide the hashcount")
|
|
settings.add_argument('-c','--column', type=int, default=1 ,help="Which Column of inputfile should be processed (0,1,...)")
|
|
settings.add_argument('-l','--label', default="MD5",help="What kind of Data is beeing processed (MD5,filenames,...)")
|
|
settings.add_argument('-d','--delimiter', default=',' , help="Which char is used to delimit columns in inputfile")
|
|
settings.add_argument('-i','--inputfile', default='/nsrl/NSRLFile.txt' , help="Path of input file")
|
|
|
|
|
|
args = parser.parse_args()
|
|
|
|
nsrl_path = args.inputfile
|
|
error_rate = args.error_rate
|
|
|
|
print "[BUILDING] Using error-rate: {}".format(error_rate)
|
|
if os.path.isfile(nsrl_path):
|
|
print "[BUILDING] Reading in NSRL Database"
|
|
if args.hashcount is None:
|
|
with open(nsrl_path) as f_line:
|
|
# Strip off header
|
|
_ = f_line.readline()
|
|
print "[BUILDING] Calculating number of hashes in NSRL..."
|
|
num_lines = sum(bl.count("\n") for bl in blocks(f_line))
|
|
else:
|
|
num_lines=args.hashcount
|
|
print "[BUILDING] There are %s hashes in the NSRL Database" % num_lines
|
|
with open(nsrl_path) as f_nsrl:
|
|
# Strip off header
|
|
_ = f_nsrl.readline()
|
|
print "[BUILDING] Creating bloomfilter"
|
|
bf = BloomFilter(num_lines, error_rate)
|
|
print "[BUILDING] Inserting hashes into bloomfilter"
|
|
# sha1 hash is in column 0
|
|
for line in f_nsrl:
|
|
hashline = line.split(args.delimiter)[args.column].strip('"')
|
|
if hashline:
|
|
try:
|
|
hash = binascii.unhexlify(hashline)
|
|
bf.add(hash)
|
|
except Exception as e:
|
|
print "[ERROR] %s" % e
|
|
print "[BUILDING] NSRL bloomfilter contains {} items.".format(len(bf))
|
|
with open('nsrl.bloom', 'wb') as nb:
|
|
bf.tofile(nb)
|
|
print "[BUILDING] Complete"
|
|
else:
|
|
print("[ERROR] No such file or directory: %s", nsrl_path)
|
|
|
|
return
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|