114e5175eSLucas Bates#!/usr/bin/env python3
214e5175eSLucas Bates
314e5175eSLucas Batesimport os
414e5175eSLucas Batesimport signal
514e5175eSLucas Batesfrom string import Template
614e5175eSLucas Batesimport subprocess
714e5175eSLucas Batesimport time
814e5175eSLucas Batesfrom TdcPlugin import TdcPlugin
914e5175eSLucas Bates
1014e5175eSLucas Batesfrom tdc_config import *
1114e5175eSLucas Bates
1214e5175eSLucas Batestry:
1314e5175eSLucas Bates    from scapy.all import *
1414e5175eSLucas Batesexcept ImportError:
1514e5175eSLucas Bates    print("Unable to import the scapy python module.")
1614e5175eSLucas Bates    print("\nIf not already installed, you may do so with:")
1714e5175eSLucas Bates    print("\t\tpip3 install scapy==2.4.2")
1814e5175eSLucas Bates    exit(1)
1914e5175eSLucas Bates
2014e5175eSLucas Batesclass SubPlugin(TdcPlugin):
2114e5175eSLucas Bates    def __init__(self):
2214e5175eSLucas Bates        self.sub_class = 'scapy/SubPlugin'
2314e5175eSLucas Bates        super().__init__()
2414e5175eSLucas Bates
2514e5175eSLucas Bates    def post_execute(self):
2614e5175eSLucas Bates        if 'scapy' not in self.args.caseinfo:
2714e5175eSLucas Bates            if self.args.verbose:
2814e5175eSLucas Bates                print('{}.post_execute: no scapy info in test case'.format(self.sub_class))
2914e5175eSLucas Bates            return
3014e5175eSLucas Bates
3114e5175eSLucas Bates        # Check for required fields
3214e5175eSLucas Bates        scapyinfo = self.args.caseinfo['scapy']
3314e5175eSLucas Bates        scapy_keys = ['iface', 'count', 'packet']
3414e5175eSLucas Bates        missing_keys = []
3514e5175eSLucas Bates        keyfail = False
3614e5175eSLucas Bates        for k in scapy_keys:
3714e5175eSLucas Bates            if k not in scapyinfo:
3814e5175eSLucas Bates                keyfail = True
3914e5175eSLucas Bates                missing_keys.add(k)
4014e5175eSLucas Bates        if keyfail:
4114e5175eSLucas Bates            print('{}: Scapy block present in the test, but is missing info:'
4214e5175eSLucas Bates                .format(self.sub_class))
4314e5175eSLucas Bates            print('{}'.format(missing_keys))
4414e5175eSLucas Bates
4514e5175eSLucas Bates        pkt = eval(scapyinfo['packet'])
4614e5175eSLucas Bates        if '$' in scapyinfo['iface']:
4714e5175eSLucas Bates            tpl = Template(scapyinfo['iface'])
4814e5175eSLucas Bates            scapyinfo['iface'] = tpl.safe_substitute(NAMES)
4914e5175eSLucas Bates        for count in range(scapyinfo['count']):
5014e5175eSLucas Bates            sendp(pkt, iface=scapyinfo['iface'])
51