def runTest(self):
import pytrap
import os
urtempl = "ipaddr IP,uint16 PORT"
c = pytrap.TrapCtx()
c.init(["-i", "f:/tmp/pytrap_test"], 0, 1)
c.setDataFmt(0, pytrap.FMT_UNIREC, urtempl)
t = pytrap.UnirecTemplate("ipaddr IP,uint16 PORT")
t.createMessage()
t.IP = pytrap.UnirecIPAddr("192.168.0.1")
t.PORT = 123
c.send(t.getData())
c.sendFlush()
c.finalize()
c = pytrap.TrapCtx()
c.init(["-i", "f:/tmp/pytrap_test"], 1)
c.setRequiredFmt(0, pytrap.FMT_UNIREC, urtempl)
data = c.recv()
t = pytrap.UnirecTemplate(urtempl)
t.setData(data)
self.assertEqual(t.IP, pytrap.UnirecIPAddr("192.168.0.1"))
self.assertEqual(t.PORT, 123)
c.finalize()
os.unlink("/tmp/pytrap_test")
def main():
# Select input spec for basic and smtp interface
basic_spec = "ipaddr DST_IP,ipaddr SRC_IP,uint64 BYTES,uint64 LINK_BIT_FIELD,time TIME_FIRST,time TIME_LAST,uint32 PACKETS,uint16 DST_PORT,uint16 SRC_PORT,uint8 DIR_BIT_FIELD,uint8 PROTOCOL,uint8 TCP_FLAGS,uint8 TOS,uint8 TTL"
smtp_spec = "ipaddr DST_IP,ipaddr SRC_IP,uint64 BYTES,uint64 LINK_BIT_FIELD,time TIME_FIRST,time TIME_LAST,uint32 PACKETS,uint32 SMTP_2XX_STAT_CODE_COUNT,uint32 SMTP_3XX_STAT_CODE_COUNT,uint32 SMTP_4XX_STAT_CODE_COUNT,uint32 SMTP_5XX_STAT_CODE_COUNT,uint32 SMTP_COMMAND_FLAGS,uint32 SMTP_MAIL_CMD_COUNT,uint32 SMTP_RCPT_CMD_COUNT,uint32 SMTP_STAT_CODE_FLAGS,uint16 DST_PORT,uint16 SRC_PORT,uint8 DIR_BIT_FIELD,uint8 PROTOCOL,uint8 TCP_FLAGS,uint8 TOS,uint8 TTL,string SMTP_DOMAIN,string SMTP_FIRST_RECIPIENT,string SMTP_FIRST_SENDER"
# Create a new trap context
trap = pytrap.TrapCtx()
"""
Trap initialization for two input interfaces, and no output interface
"""
trap.init(sys.argv, 2, 1)
# Set up required format to accept any unirec format.
trap.setRequiredFmt(BASIC_IF, pytrap.FMT_UNIREC,
basic_spec) # Refers to flows without SMTP headers
trap.setRequiredFmt(SMTP_IF, pytrap.FMT_UNIREC,
smtp_spec) # Refers to flows with SMTP headers
basic_rec = pytrap.UnirecTemplate(basic_spec)
smtp_rec = pytrap.UnirecTemplate(smtp_spec)
trap.setDataFmt(0, pytrap.FMT_JSON, "smtp_alert")
log.info("Daemon: Trap initialized.")
detector = SpamDetection(trap)
# Data synchronized queues
flow_queue = Queue() # Synchronize input flows
# Create workers for each receiver
basic_rcv = Thread(name="basic_receiver",
target=fetch_data,
args=(trap, basic_rec, BASIC_IF, flow_queue))
smtp_rcv = Thread(name="smtp_receiver",
target=fetch_data,
args=(trap, smtp_rec, SMTP_IF, flow_queue))
# Handle the received data from receivers
data_handler = Thread(name="data_handler",
target=data_handling,
args=(detector, flow_queue))
# Run multi-receiver
basic_rcv.start()
smtp_rcv.start()
data_handler.start()
log.info("Daemon: Multi-receiver started.")
# Start detector
detector.start()
# Wait until the daemon is requested to stop by releasing the lock (by signal handler)
g.stop_lock.acquire()
signal.signal(signal.SIGINT, signal.SIG_DFL)
signal.signal(signal.SIGTERM, signal.SIG_DFL)
signal.signal(signal.SIGABRT, signal.SIG_DFL)
log.info("Stopping running components ...")
g.is_running = False
detector.stop()
# Join the threads
basic_rcv.join()
smtp_rcv.join()
# Stop data_handler
flow_queue.put(None)
data_handler.join()
detector.join()
# Free allocated memory
trap.finalize()
log.info("***** Finished, main thread exiting. *****")
logging.shutdown()
sys.exit(0)
def runTest(self):
import pytrap
c = pytrap.TrapCtx()
try:
c.getInIFCState(0)
self.fail("Calling method of uninitialized context.")
except:
pass
def runTest(self):
import pytrap
c = pytrap.TrapCtx()
try:
c.setRequiredFmt(0)
self.fail("Calling method of uninitialized context.")
except:
pass
def runTest(self):
import pytrap
c = pytrap.TrapCtx()
try:
c.ifcctl(0, True, CTL_IMEOUT, 0)
self.fail("Calling method of uninitialized context.")
except:
pass
def runTest(self):
import pytrap
c = pytrap.TrapCtx()
try:
c.init(["-h", "1"], 0, 1)
self.fail("Calling method of uninitialized context.")
except pytrap.TrapHelp:
pass
c.finalize()
def __init__(self, interfaces, callback):
unirec_init = ["-i", ",".join(x["interface"] for x in interfaces)]
self.interfaces = interfaces
self.callback = callback
self.threads = []
self.lock = threading.Lock()
self.trap = pytrap.TrapCtx()
self.trap.init(unirec_init, len(interfaces), 0)
for i in range(len(interfaces)):
self.trap.setRequiredFmt(i, pytrap.FMT_UNIREC,
"ipaddr SRC_IP,ipaddr DST_IP")
def runTest(self):
#"""json.dump returns str object, which was formerly not supported by pytrep send()"""
import pytrap
import json
import os
urtempl = "test"
jdata = {"SRC_IP": "192.168.0.1"}
c = pytrap.TrapCtx()
c.init(["-i", "f:/tmp/pytrap_test2"], 0, 1)
c.setDataFmt(0, pytrap.FMT_JSON, urtempl)
c.send(json.dumps(jdata))
c.sendFlush()
c.finalize()
c = pytrap.TrapCtx()
c.init(["-i", "f:/tmp/pytrap_test2"], 1)
c.setRequiredFmt(0, pytrap.FMT_JSON, urtempl)
data = c.recv()
c.finalize()
self.assertEqual(jdata, json.loads(data))
os.unlink("/tmp/pytrap_test2")
def load_pytrap():
"""Init nemea libraries and set format of IP flows.
Returns
--------
rec : pytrap
Templete of IP flows.
trap : pytrap
Init pytrap NEMEA library, for capture IP flows from IFC interface setted in paramater -i.
"""
trap = pytrap.TrapCtx()
trap.init(sys.argv)
# Set the list of required fields in received messages.
# This list is an output of e.g. flow_meter - basic flow.
input_spec = "ipaddr DST_IP,ipaddr SRC_IP,uint64 BYTES,uint64 LINK_BIT_FIELD,time TIME_FIRST,time TIME_LAST,uint32 PACKETS,uint16 DST_PORT,uint16 SRC_PORT,uint8 DIR_BIT_FIELD,uint8 PROTOCOL,uint8 TCP_FLAGS"
trap.setRequiredFmt(0, pytrap.FMT_UNIREC, input_spec)
rec = pytrap.UnirecTemplate(input_spec)
return rec, trap
def runTest(self):
import pytrap
c = pytrap.TrapCtx()
try:
c.init()
self.fail("argv is mandatory")
except:
pass
try:
c.init([])
self.fail("argv must not be empty")
except:
pass
try:
c.init([[]])
self.fail("argv[i] must be string")
except:
pass
c.init(["-i", "u:test_init"], 0, 1)
c.finalize()
def amendflows(args, resolver):
"""Read flow data, add resolved fields and send amended flow data."""
trap = pytrap.TrapCtx()
trap.init(sys.argv, 1, 1)
urinput = pytrap.UnirecTemplate(args.urformat)
# this module accepts all Unirec fieds -> set required format:
trap.setRequiredFmt(0, pytrap.FMT_UNIREC, args.urformat)
trap.ifcctl(0, False, pytrap.CTL_BUFFERSWITCH, 0)
uroutformat = urspecaddresolvedfields(args.urformat, args.outfields)
uroutput = pytrap.UnirecTemplate(uroutformat)
trap.setDataFmt(0, pytrap.FMT_UNIREC, uroutformat)
if args.verbose:
print("Set output format to '{}'."
"\nWaiting for events.".format(uroutformat))
while True:
# Read data from input interface
try:
indata = trap.recv()
if args.verbose:
sys.stdout.write('.')
except pytrap.FormatMismatch:
print("Error: input and output interfaces data format"
" or data specifier mismatch")
break
except pytrap.FormatChanged as err:
# Get data format from negotiation, amend it, and set it
# for output iface
fmttype, fmtspec = trap.getDataFmt(0) # pylint:disable=unused-variable
# Update Unirec templates
urinput = pytrap.UnirecTemplate(fmtspec)
uroutformat = urspecaddresolvedfields(fmtspec, args.outfields)
uroutput = pytrap.UnirecTemplate(uroutformat)
trap.setDataFmt(0, pytrap.FMT_UNIREC, uroutformat)
if args.verbose:
print("Reset output format to '{}'.".format(uroutformat))
# Store data from the exception
indata = err.data # pylint:disable=no-member
except pytrap.Terminated:
print("Terminated trap.")
break
except pytrap.TrapError:
break
# Check for "end-of-stream" record
if len(indata) <= 1:
break
# Set indata for access using attributes
urinput.setData(indata)
inputvarlensize = urinput.recVarlenSize()
# Copy flow info from indata to outdata
resolvedvarlenmaxsize = sum(ftype[1] for fname, ftype in args.outfields
if len(ftype) > 1)
outdata = uroutput.createMessage(inputvarlensize +
resolvedvarlenmaxsize)
uroutput.setData(outdata)
for attr, value in urinput:
setattr(uroutput, attr, value)
for fieldstoresolve, resolution, outfield in args.validatedresolvspecs:
setattr(uroutput, outfield,
getattr(resolver, resolution)(urinput, *fieldstoresolve))
try:
trap.send(outdata)
except pytrap.Terminated:
print("Terminated trap.")
break
if args.verbose:
sys.stdout.write(',')
trap.sendFlush()
"--threshold",
default=5000,
help="set number of JSON records per file")
parser.add_argument('-d',
"--destination",
default='./',
help="target output directory.")
parser.add_argument(
"--tmp-dir",
default='./',
help=
"directory for storing temporary file before moving them to output destination."
)
args = parser.parse_args()
context = pytrap.TrapCtx()
context.init(['-i', args.ifcspec])
context.setRequiredFmt(0, pytrap.FMT_UNIREC, "")
unirec = pytrap.UnirecTemplate("ipaddr DST_IP,ipaddr SRC_IP,uint64 BYTES,uint64 LINK_BIT_FIELD,time TIME_FIRST, "\
"time TIME_LAST,uint32 DNS_RR_TTL,uint32 PACKETS,uint16 DNS_ANSWERS,uint16 DNS_CLASS,"\
"uint16 DNS_ID,uint16 DNS_PSIZE,uint16 DNS_QTYPE,uint16 DNS_RLENGTH,uint16 DST_PORT,"\
"uint16 SRC_PORT,uint8 DIR_BIT_FIELD,uint8 DNS_DO,uint8 DNS_RCODE,"\
"uint8 PROTOCOL,uint8 TCP_FLAGS,uint8 TOS,uint8 TTL,string DNS_NAME,bytes DNS_RDATA")
tmp_dir = args.tmp_dir
if not os.path.exists(tmp_dir):
sys.stderr.write("Path given by --tmp-dir not found.\n")
exit(1)
output_path = args.destination
def Run(module_name,
module_desc,
req_type,
req_format,
conv_func,
arg_parser=None):
""" TODO doc
"""
# *** Parse command-line arguments ***
if arg_parser is None:
arg_parser = argparse.ArgumentParser()
arg_parser.formatter_class = argparse.RawDescriptionHelpFormatter
# Set description
arg_parser.description = str.format(
desc_template,
name=module_name,
type={
pytrap.FMT_RAW: 'raw',
pytrap.FMT_UNIREC: 'UniRec',
pytrap.FMT_JSON: 'JSON'
}.get(req_type, '???'),
fmt=req_format,
original_desc=module_desc + "\n\n " if module_desc else "",
)
# Add arguments defining outputs
# TRAP output
arg_parser.add_argument(
'--trap',
action='store_true',
help=
'Enable output via TRAP interface (JSON type with format id "IDEA"). Parameters are set using "-i" option as usual.'
)
# File output
arg_parser.add_argument(
'--file',
metavar="FILE",
type=str,
help=
'Enable output to file (each IDEA message printed to new line in JSON format). Set to "-" to use standard output.'
)
arg_parser.add_argument(
'--file-indent',
metavar="N",
type=int,
help='Pretty-format JSON in output file using N spaces for indentation.'
)
arg_parser.add_argument('--file-append',
action='store_true',
help='Append to file instead of overwrite.')
# MongoDB output
arg_parser.add_argument(
'--mongodb',
metavar="DBNAME",
help='Enable output to MongoDB. Connect to database named DBNAME.')
arg_parser.add_argument(
'--mongodb-coll',
metavar="COLL",
default='alerts',
help='Put IDEA messages into collection named COLL (default: "alerts").'
)
arg_parser.add_argument(
'--mongodb-host',
metavar="HOSTNAME",
default='localhost',
help='Connect to MongoDB running on HOSTNAME (default: "localhost").')
arg_parser.add_argument(
'--mongodb-port',
metavar="PORT",
type=int,
default=27017,
help='Connect to MongoDB running on port number PORT (default: 27017).'
)
# Warden3 output
arg_parser.add_argument(
'--warden',
metavar="CONFIG_FILE",
help=
'Send IDEA messages to Warden server. Load configuration of Warden client from CONFIG_FILE.'
)
# Other options
arg_parser.add_argument(
'-n',
'--name',
metavar='NODE_NAME',
help=
'Name of the node, filled into "Node.Name" element of the IDEA message. Required if Warden output is used, recommended otherwise.'
)
arg_parser.add_argument(
'--test',
action='store_true',
help='Add "Test" to "Category" before sending a message to output(s).')
arg_parser.add_argument(
'-v',
'--verbose',
action='store_true',
help=
"Enable verbose mode (may be used by some modules, common part donesn't print anything"
)
arg_parser.add_argument(
'--srcwhitelist-file',
metavar="FILE",
type=str,
help=
"File with addresses/subnets in format: <ip address>/<mask>,<data>\\n \n where /<mask>,<data> is optional, <data> is a user-specific optional content. Whitelist is applied to SRC_IP field. If SRC_IP from the alert is on whitelist, the alert IS NOT reported."
)
arg_parser.add_argument(
'--dstwhitelist-file',
metavar="FILE",
type=str,
help=
"File with addresses/subnets, whitelist is applied on DST_IP, see --srcwhitelist-file help."
)
# TRAP parameters
trap_args = arg_parser.add_argument_group('Common TRAP parameters')
trap_args.add_argument(
'-i',
metavar="IFC_SPEC",
required=True,
help='TODO (ideally this section should be added by TRAP')
# Parse arguments
args = arg_parser.parse_args()
# Check if at least one output is enabled
if not (args.file or args.trap or args.mongodb or args.warden):
sys.stderr.write(module_name +
": Error: At least one output must be selected\n")
exit(1)
# Check if node name is set if Warden output is enabled
if args.name is None:
if args.warden:
sys.stderr.write(
module_name +
": Error: Node name must be specified if Warden output is used (set param --name).\n"
)
exit(1)
else:
sys.stderr.write(module_name +
": Warning: Node name is not specified.\n")
# *** Initialize TRAP ***
trap = pytrap.TrapCtx()
trap.init(["-i", args.i], 1, 1 if args.trap else 0)
#trap.setVerboseLevel(3)
#trap.registerDefaultSignalHandler()
# Set required input format
trap.setRequiredFmt(0, req_type, req_format)
# If TRAP output is enabled, set output format (JSON, format id "IDEA")
if args.trap:
trap.setDataFmt(0, pytrap.FMT_JSON, "IDEA")
# *** Create output handles/clients/etc ***
filehandle = None
mongoclient = None
mongocoll = None
wardenclient = None
if args.file:
if args.file == '-':
filehandle = sys.stdout
else:
filehandle = open(args.file, "a" if args.file_append else "w")
if args.mongodb:
import pymongo
mongoclient = pymongo.MongoClient(args.mongodb_host, args.mongodb_port)
mongocoll = mongoclient[args.mongodb][args.mongodb_coll]
if args.warden:
import warden_client
try:
config = warden_client.read_cfg(args.warden)
except ValueError as e:
sys.stderr.write(
"{0}: Failed to load Warden config file '{1}'\n{2}\n".format(
module_name, args.warden, e))
exit(1)
config['name'] = args.name
wardenclient = warden_client.Client(**config)
# Check if a whitelist is set, parse the file and prepare context for binary search
import ip_prefix_search
if args.srcwhitelist_file:
if 'ipaddr SRC_IP' in req_format.split(","):
srcwhitelist = ip_prefix_search.IPPSContext.fromFile(
args.srcwhitelist_file)
else:
srcwhitelist = None
if args.dstwhitelist_file:
if 'ipaddr DST_IP' in req_format.split(","):
dstwhitelist = ip_prefix_search.IPPSContext.fromFile(
args.dstwhitelist_file)
else:
dstwhitelist = None
# *** Main loop ***
URInputTmplt = None
if req_type == pytrap.FMT_UNIREC and req_format != "":
URInputTmplt = pytrap.UnirecTemplate(
req_format
) # TRAP expects us to have predefined template for required set of fields
rec = URInputTmplt
stop = False
while not stop:
# *** Read data from input interface ***
try:
data = trap.recv()
except pytrap.FormatMismatch:
sys.stderr.write(
module_name + ": Error: input data format mismatch\n"
) #Required: "+str((req_type,req_format))+"\nReceived: "+str(trap.get_data_fmt(trap.IFC_INPUT, 0))+"\n")
break
except pytrap.FormatChanged as e:
# TODO: This should be handled by trap.recv transparently
# Get negotiated input data format
(fmttype, fmtspec) = trap.getDataFmt(0)
# If data type is UniRec, create UniRec template
if fmttype == pytrap.FMT_UNIREC:
URInputTmplt = pytrap.UnirecTemplate(fmtspec)
else:
URInputTmplt = None
rec = URInputTmplt
data = e.data
except pytrap.Terminated:
break
# Check for "end-of-stream" record
if len(data) <= 1:
# If we have output, send "end-of-stream" record and exit
if args.trap:
trap.send(0, b"0")
break
# Assert that if UniRec input is required, input template is set
assert (req_type != pytrap.FMT_UNIREC or URInputTmplt is not None)
# Convert raw input data to UniRec object (if UniRec input is expected)
if req_type == pytrap.FMT_UNIREC:
rec.setData(data)
elif req_type == pytrap.FMT_JSON:
rec = json.loads(data)
else: # TRAP_FMT_RAW
rec = data
# Check whitelists
if srcwhitelist and srcwhitelist.ip_search(rec.SRC_IP):
continue
if dstwhitelist and dstwhitelist.ip_search(rec.DST_IP):
continue
# *** Convert input record to IDEA ***
# Pass the input record to conversion function to create IDEA message
idea = conv_func(rec, args)
if idea is None:
continue # Record can't be converted - skip it (notice should be printed by the conv function)
if args.name is not None:
idea['Node'][0]['Name'] = args.name
if args.test:
idea['Category'].append('Test')
# *** Send IDEA to outputs ***
# File output
if filehandle:
filehandle.write(json.dumps(idea, indent=args.file_indent) + '\n')
# TRAP output
if args.trap:
try:
trap.send(json.dumps(idea), 0)
except pytrap.TimeoutError:
# skip this message
pass
except pytrap.Terminated:
# don't exit immediately, first finish sending to other outputs
stop = True
# MongoDB output
if mongocoll:
# We need to change IDEA message here, but we may need it unchanged
# later -> copy it (shallow copy is sufficient)
idea2 = idea.copy()
# Convert timestamps from string to Date format
idea2['DetectTime'] = datetime.strptime(idea2['DetectTime'],
"%Y-%m-%dT%H:%M:%SZ")
for i in ['CreateTime', 'EventTime', 'CeaseTime']:
if idea2.has_key(i):
idea2[i] = datetime.strptime(idea2[i],
"%Y-%m-%dT%H:%M:%SZ")
try:
mongocoll.insert(idea2)
except pymongo.errors.AutoReconnect:
sys.stderr.write(module_name +
": Error: MongoDB connection failure.\n")
stop = True
# Warden output
if wardenclient:
wardenclient.sendEvents([idea])
# *** Cleanup ***
if filehandle and filehandle != sys.stdout:
filehandle.close()
if mongoclient:
mongoclient.close()
if wardenclient:
wardenclient.close()
trap.finalize()
def Run(module_name,
module_desc,
req_type,
req_format,
conv_func,
arg_parser=None):
"""Run the main loop of the reporter module called `module_name` with `module_desc` (used in help).
The module requires data format of `req_type` type and `req_format` specifier - these must be given by author of the module.
`conv_func(rec, args)` is a callback function that must translate given incoming alert `rec` (typically in UniRec according to `req_type`) into IDEA message. `args` contains CLI arguments parsed by ArgumentParser. `conv_func` must return dict().
"""
# *** Parse command-line arguments ***
if arg_parser is None:
arg_parser = argparse.ArgumentParser()
arg_parser.formatter_class = argparse.RawDescriptionHelpFormatter
# Set description
arg_parser.description = str.format(
desc_template,
name=module_name,
type={
pytrap.FMT_RAW: 'raw',
pytrap.FMT_UNIREC: 'UniRec',
pytrap.FMT_JSON: 'JSON'
}.get(req_type, '???'),
fmt=req_format,
original_desc=module_desc + "\n\n " if module_desc else "",
)
# Add arguments defining outputs
# TRAP output
arg_parser.add_argument(
'-T',
'--trap',
action='store_true',
help=
'Enable output via TRAP interface (JSON type with format id "IDEA"). Parameters are set using "-i" option as usual.'
)
# Config file
arg_parser.add_argument(
'-c',
'--config',
metavar="FILE",
default="./config.yaml",
type=str,
help='Specify YAML config file path which to load.')
arg_parser.add_argument('-d',
'--dry',
action='store_true',
help="""Do not run, just print loaded config.""")
# Warden3 output
arg_parser.add_argument(
'-W',
'--warden',
metavar="CONFIG_FILE",
help=
'Send IDEA messages to Warden server. Load configuration of Warden client from CONFIG_FILE.'
)
# Other options
arg_parser.add_argument(
'-n',
'--name',
metavar='NODE_NAME',
help=
'Name of the node, filled into "Node.Name" element of the IDEA message. Required argument.'
)
arg_parser.add_argument(
'-v',
'--verbose',
metavar='VERBOSE_LEVEL',
default=3,
type=int,
help=
"""Enable verbose mode (may be used by some modules, common part doesn't print anything).\nLevel 1 logs everything, level 5 only critical errors. Level 0 doesn't log."""
)
# TRAP parameters
trap_args = arg_parser.add_argument_group('Common TRAP parameters')
trap_args.add_argument(
'-i',
metavar="IFC_SPEC",
required=True,
help=
'See http://nemea.liberouter.org/trap-ifcspec/ for more information.')
# Parse arguments
args = arg_parser.parse_args()
# Set log level
logging.basicConfig(level=(args.verbose * 10), format=FORMAT)
# Check if node name is set if Warden output is enabled
if args.name is None:
#if args.warden:
# sys.stderr.write(module_name+": Error: Node name must be specified if Warden output is used (set param --name).\n")
# exit(1)
logger.warning("Node name is not specified.")
# *** Initialize TRAP ***
trap = pytrap.TrapCtx()
logger.info("Trap arguments: %s", args.i)
trap.init(["-i", args.i], 1, 1 if args.trap else 0)
#trap.setVerboseLevel(3)
#trap.registerDefaultSignalHandler()
# Set required input format
trap.setRequiredFmt(0, req_type, req_format)
if args.trap:
trap.setDataFmt(0, pytrap.FMT_JSON, "IDEA")
# *** Create output handles/clients/etc ***
wardenclient = None
if args.warden:
try:
import warden_client
except:
logger.error(
"There is no available warden_client python module. Install it or remove '--warden' from the module's arguments."
)
sys.exit(1)
config = warden_client.read_cfg(args.warden)
config['name'] = args.name
wardenclient = warden_client.Client(**config)
# Initialize configuration
config = Config.Config(args.config, trap=trap, warden=wardenclient)
if not args.dry:
# *** Main loop ***
URInputTmplt = None
if req_type == pytrap.FMT_UNIREC and req_format != "":
URInputTmplt = pytrap.UnirecTemplate(
req_format
) # TRAP expects us to have predefined template for required set of fields
rec = URInputTmplt
stop = False
while not stop:
logger.info("Starting receiving")
# *** Read data from input interface ***
try:
data = trap.recv()
except pytrap.FormatMismatch:
logger.error(
"Input data format mismatch in receiving from TRAP interface"
)
break
except pytrap.FormatChanged as e:
# Get negotiated input data format
(fmttype, fmtspec) = trap.getDataFmt(0)
# If data type is UniRec, create UniRec template
if fmttype == pytrap.FMT_UNIREC:
URInputTmplt = pytrap.UnirecTemplate(fmtspec)
else:
URInputTmplt = None
rec = URInputTmplt
data = e.data
except pytrap.Terminated:
break
# Check for "end-of-stream" record
if len(data) <= 1:
if args.trap:
# If we have output, send "end-of-stream" record and exit
trap.send(0, b"0")
break
# Assert that if UniRec input is required, input template is set
assert (req_type != pytrap.FMT_UNIREC or URInputTmplt is not None)
# Convert raw input data to UniRec object (if UniRec input is expected)
if req_type == pytrap.FMT_UNIREC:
rec.setData(data)
elif req_type == pytrap.FMT_JSON:
rec = json.loads(data)
else: # TRAP_FMT_RAW
rec = data
# *** Convert input record to IDEA ***
# Pass the input record to conversion function to create IDEA message
idea = conv_func(rec, args)
if idea is None:
# Record can't be converted - skip it
continue
if args.name is not None:
idea['Node'][0]['Name'] = args.name
# *** Send IDEA to outputs ***
# Perform rule matching and action running on the idea message
try:
config.match(idea)
except pytrap.Terminated:
logger.error("PyTrap was terminated")
break
except DropMsg:
logger.info("Message was dropped by Drop action.")
continue
except Exception as e:
logger.error(str(e))
break
else:
# DRY argument given, just print config and exit
print(config)
if wardenclient:
wardenclient.close()
trap.finalize()
def runTest(self):
import pytrap
import os
import time
messages = 10000000
urtempl = "ipaddr IP,uint16 PORT"
# Start sender
c1 = pytrap.TrapCtx()
c1.init(["-i", "f:/tmp/pytrap_test3"], 0, 1)
c1.setDataFmt(0, pytrap.FMT_UNIREC, urtempl)
c1.ifcctl(0, False, pytrap.CTL_TIMEOUT, 500000)
c1.ifcctl(0, False, pytrap.CTL_AUTOFLUSH, 500000)
t = pytrap.UnirecTemplate(urtempl)
t.createMessage()
t.IP = pytrap.UnirecIPAddr("192.168.0.1")
for i in range(messages):
t.PORT = i
c1.send(t.getData())
c1.sendFlush()
# Start Receiver
c2 = pytrap.TrapCtx()
c2.init(["-i", "f:/tmp/pytrap_test3"], 1)
c2.setRequiredFmt(0, pytrap.FMT_UNIREC, urtempl)
startt = time.process_time()
data = c2.recvBulk(t, time=15, count=messages)
elapsed_time = time.process_time() - startt
print(
f"recvBulk() Elapsed time for {messages} messages is: {elapsed_time}"
)
self.assertEqual(len(data), messages)
ports = [i["PORT"] for i in data]
self.assertEqual(ports, [i & 0xFFFF for i in range(messages)])
# Start Receiver
c2 = pytrap.TrapCtx()
c2.init(["-i", "f:/tmp/pytrap_test3"], 1)
c2.setRequiredFmt(0, pytrap.FMT_UNIREC, urtempl)
startt = time.process_time()
data = list()
while True:
d = c2.recv()
if not d:
break
t.setData(d)
data.append(t.getDict())
elapsed_time = time.process_time() - startt
print(
f"recv() Elapsed time for {messages} messages is: {elapsed_time}")
self.assertEqual(len(data), messages)
ports = [i["PORT"] for i in data]
self.assertEqual(ports, [i & 0xFFFF for i in range(messages)])
c1.finalize()
c2.finalize()
os.unlink("/tmp/pytrap_test3")
send_to_reporter(detection_event)
def tear_down(self):
self.receiver.join_and_quit()
self.processor.stop()
if __name__ == '__main__':
options, args = parser.parse_args()
signal.signal(signal.SIGINT, signal_h)
g.blacklists = utils.load_blacklists(options.blacklist_config)
g.botnet_blacklist_indexes = utils.get_botnet_blacklist_indexes(
g.blacklists)
logger.setLevel(options.log_level)
trap = pytrap.TrapCtx()
trap.init(sys.argv, 2, 2)
trap.setDataFmt(0, pytrap.FMT_JSON, "aggregated_blacklist")
trap.setDataFmt(1, pytrap.FMT_JSON, "blacklist_evidence")
controller = Controller(options.adaptive_blacklist,
options.process_interval, options.evidence_timeout)
# Run Controller until stopped by signal
controller.run()
controller.tear_down()
# Free allocated memory
trap.finalize()
#!/usr/bin/env python3
import pytrap
import unirec
import sys
import pdb
#UR_Flow = unirec.CreateTemplate("UR_Flow", b"ipaddr DST_IP,ipaddr SRC_IP,double EVENT_SCALE,time TIME_FIRST,time TIME_LAST,uint16 DST_PORT,uint16 SRC_PORT,uint8 EVENT_TYPE,uint8 PROTOCOL")
UR_Flow = unirec.CreateTemplate("UR_Flow", b"double EVENT_SCALE,time TIME_FIRST,time TIME_LAST,uint16 DST_PORT,uint16 SRC_PORT,uint8 EVENT_TYPE,uint8 PROTOCOL")
ctx = pytrap.TrapCtx()
ctx.init(sys.argv)
#pytrap.setVerboseLevel(2)
print(ctx.getVerboseLevel())
ctx.setRequiredFmt(0)
ctx.ifcctl(0, True, 3, 1000000)
try:
a = ctx.recv(0)
raise Exception("recv should have raised Timeout exception")
except pytrap.TrapError as e:
print("Caught Timeout exception: {0}".format(type(e)))
ctx.finalize()
def main():
# Process arguments and initialize TRAP interface
try:
args, _ = process_arguments()
except SystemExit as e:
args = None # Error in arguments or help message
trap = pytrap.TrapCtx()
try:
trap.init(sys.argv,
1,
0,
module_name="Backscatter classifier common TRAP help")
# Allow trap to print it's own help but terminate script anyway due to error in arguments
if args is None:
sys.exit(PYTHON_ARGS_ERROR_EXIT)
trap.setRequiredFmt(0)
except Exception as e:
# Trap error message
print(e)
sys.exit(TRAP_ARGS_ERROR_EXIT)
# Logging settings
logger = logging.getLogger("backscatter_classifier")
logging.basicConfig(
level=logging.DEBUG,
filename=args.logfile,
filemode='w',
format=
"[%(levelname)s], %(asctime)s, %(name)s, %(funcName)s, line %(lineno)d: %(message)s"
)
# ASN and city databases
try:
geoip_db = Geoip2Wrapper(args.agp, args.cgp)
except Exception as e:
logger.error(e)
logger.error("Error while create GeoIP2 wrapper")
print(str(e), file=sys.stderr)
sys.exit(EXIT_FAILURE)
if args.export_to in ("c3isp", "c3isp-misp"):
c3isp_upload.read_config(args.c3isp_config)
if args.export_to == "misp":
# MISP instance
try:
misp_instance = ExpandedPyMISP(args.url, args.key, args.ssl)
except Exception as e:
logger.error(e)
logger.error("Error while creating MISP instance")
print(str(e), file=sys.stderr)
sys.exit(EXIT_FAILURE)
# DDoS model
ddos_model = pickle.load(args.model)
ddos_classifier = DDoSClassifier(ddos_model, args.min_threshold)
# *** MAIN PROCESSING LOOP ***
while True:
try:
# Receive data
try:
data = trap.recv()
except pytrap.FormatChanged as e:
# Automatically detect format
fmttype, fmtspec = trap.getDataFmt(0)
rec = pytrap.UnirecTemplate(fmtspec)
data = e.data
if len(data) <= 1:
# Terminating message
break
# Decode received data into python object
rec.setData(data)
# Only Ipv4
try:
victim_ip = IPv4Address(rec.SRC_IP)
except Exception as e:
logger.info("Received IPv6 address, skipping")
continue
# Predict class of backscatter like traffic
try:
duration = DDoSClassifier.DURATION(rec)
if duration < args.min_duration:
# attack is too short lived to be reported
continue
if duration > args.max_duration:
# attack is too long
continue
if rec.FLOW_COUNT < args.min_flows:
# attack does not have enough flows
continue
for subnet in CESNET_NET:
if victim_ip in subnet:
continue
ddos = ddos_classifier.predict(rec)
except Exception as e:
logger.error(e)
continue
# Report attack using MISP
try:
if ddos:
try:
domain = gethostbyaddr(str(victim_ip))
except herror as e:
# Do not report for unknown domains
continue
event = create_ddos_event(rec, geoip_db, victim_ip,
domain[0],
args.misp_templates_dir)
if args.export_to == "misp":
try:
event_id = misp_instance.add_event(
event)['Event']['id']
misp_instance.publish(event_id)
except Exception as e:
logger.error(e)
elif args.export_to in ("c3isp", "c3isp-misp"):
logger.debug(f"Uploading event to C3ISP platform.")
event_file_path = Path(
"misp_event_to_c3isp.json").absolute()
with temporary_open(event_file_path,
'w') as event_file:
json.dump(event.to_json(), event_file)
response = c3isp_upload.upload_to_c3isp(
event_file_path)
logger.debug(f"Response: {response}")
if not response or ('status' in response
and response['status'] == 'ERROR'):
logger.error("ERROR during upload!")
continue
if args.export_to == "c3isp-misp":
dpo_id = response['content'][
'additionalProperties']['dposId']
logger.debug(f'Exporting DPO {dpo_id} to MISP.')
c3isp_upload.export_misp(dpo_id, logger)
except Exception as e:
logger.error(str(e))
continue
except Exception as e:
# Log and re-raise exception
logger.error(e)
raise e
# *** END OF MAIN PROCESSING LOOP ***
trap.finalize()
