def _unset_filter(self, mark, eth):
"""Given a mark and an interface, delete the filter.
Args:
mark: The mark based on which we delete the filter.
eth: The interface on which we delete the filter.
Returns:
A TrafficControlRc containing information on success/failure.
"""
ifid = eth['id']
parent = 0x10000
self.logger.info("deleting filter on IFID {0}, handle {1:X}".format(
eth['name'], mark))
try:
self.ipr.tc(RTM_DELTFILTER,
'fw',
ifid,
mark,
parent=parent,
protocol=ETH_P_IP,
prio=PRIO)
except NetlinkError as e:
return TrafficControlRc(code=ReturnCode.NETLINK_FW_ERROR,
message=str(e))
except Exception as e:
self.logger.exception('_unset_filter')
exc_info = sys.exc_info()
return TrafficControlRc(code=ReturnCode.UNKNOWN_FW_ERROR,
message=str(exc_info))
return TrafficControlRc(code=ReturnCode.OK)
def _unset_htb_class(self, mark, eth):
"""Given a mark and an interface, unset the HTB class.
Args:
mark: The mark based on which we delete the class.
eth: The interface on which to delete that class id.
Returns:
A TrafficControlRc containing information on success/failure.
"""
ifid = eth['id']
idx = 0x10000 + mark
try:
self.logger.info("deleting class on IFID {0}, classid {1}".format(
eth['name'], int_to_classid(idx)))
self.ipr.tc(RTM_DELTCLASS, 'htb', ifid, idx)
except NetlinkError as e:
return TrafficControlRc(code=ReturnCode.NETLINK_HTB_ERROR,
message=str(e))
except Exception as e:
self.logger.exception('_unset_htb_class')
exc_info = sys.exc_info()
return TrafficControlRc(code=ReturnCode.UNKNOWN_HTB_ERROR,
message=str(exc_info))
return TrafficControlRc(code=ReturnCode.OK)
def getCapturePackets(self, mac, mc):
self.logger.info("getCapturePackets mac:{0}".format(mac))
pkts = self.scapy_task.getCapturePackets(mc['ip'])
if pkts is None:
return TrafficControlRc(code=ReturnCode.CAPTURE_NOT_READY,
message="capture is not ready")
packet_dump = PacketsToJson(pkts)
# self.logger.info("packets: {}".format(packet_dump))
return TrafficControlRc(code=ReturnCode.OK, message=packet_dump)
def exportPcap(self, mac, mc):
self.logger.info("exportPcap mac:{0}".format(mac))
pkts = self.scapy_task.getCapturePackets(mc['ip'])
if pkts is None:
return TrafficControlRc(code=ReturnCode.CAPTURE_NOT_READY,
message="capture is not ready")
if not os.path.exists(self.DEFAULT_PCAP_PATH):
os.makedirs(self.DEFAULT_PCAP_PATH)
pcap_path = self.DEFAULT_PCAP_PATH + "[" + mac + "].pcap"
wrpcap(pcap_path, pkts)
return TrafficControlRc(code=ReturnCode.OK)
def _set_netem_qdisc(self, mark, eth, shaping):
"""Given a mark, interface and shaping settings, create the NetEm
Qdisc.
Args:
mark: The mark based on which we create the Qdisc.
eth: The interface on which we will create the Qdisc.
shaping: The shaping settings for that interface.
Returns:
A TrafficControlRc containing information on success/failure.
"""
ifid = eth['id']
parent = 0x10000 + mark
idx = 0 # automatically assign a handleid
self.logger.info("create new Netem qdisc on IFID {0}, parent {1},"
" loss {2}%, delay {3}".format(
eth['name'], int_to_classid(parent),
shaping.loss.percentage,
shaping.delay.delay * 1000))
try:
self.ipr.tc(
RTM_NEWQDISC,
'netem',
ifid,
idx,
parent=parent,
loss=shaping.loss.percentage,
delay=shaping.delay.delay * 1000,
jitter=shaping.delay.jitter * 1000,
delay_corr=shaping.delay.correlation,
loss_corr=shaping.loss.correlation,
prob_reorder=shaping.reorder.percentage,
corr_reorder=shaping.reorder.correlation,
gap=shaping.reorder.gap,
prob_corrupt=shaping.corruption.percentage,
corr_corrupt=shaping.corruption.correlation,
)
except NetlinkError as e:
return TrafficControlRc(code=ReturnCode.NETLINK_NETEM_ERROR,
message=str(e))
except Exception as e:
self.logger.exception('_set_netem_qdisc')
exc_info = sys.exc_info()
return TrafficControlRc(code=ReturnCode.UNKNOWN_NETEM_ERROR,
message=str(exc_info))
return TrafficControlRc(code=ReturnCode.OK)
def _set_filter(self, mark, eth, shaping):
"""Given a mark, interface and shaping settings, create a TC filter.
Args:
mark: The mark based on which we create the filter.
eth: The interface on which we create the filter.
shaping: The shaping associated to this interface.
Returns:
A TrafficControlRc containing information on success/failure.
"""
ifid = eth['id']
idx = 0x10000 + mark
parent = 0x10000
self.logger.info("create new FW filter on IFID {0}, classid {1},"
" handle {2:X}, rate: {3}kbits".format(
eth['name'], int_to_classid(idx), mark,
shaping.rate))
try:
extra_args = {}
if not self.dont_drop_packets:
extra_args.update({
'rate':
"{}kbit".format(shaping.rate or 2**22 - 1),
'burst':
self.burst_size,
'action':
'drop',
})
self.ipr.tc(RTM_NEWTFILTER,
'fw',
ifid,
mark,
parent=parent,
protocol=ETH_P_IP,
prio=PRIO,
classid=idx,
**extra_args)
except NetlinkError as e:
return TrafficControlRc(code=ReturnCode.NETLINK_FW_ERROR,
message=str(e))
except Exception as e:
self.logger.exception('_set_filter')
exc_info = sys.exc_info()
return TrafficControlRc(code=ReturnCode.UNKNOWN_FW_ERROR,
message=str(exc_info))
return TrafficControlRc(code=ReturnCode.OK)
def decorator(self, mac, *args, **kwargs):
if not self._machine_controls.get(mac):
# can not shape while not exist
return TrafficControlRc(
code=ReturnCode.INVALID_ADDRESS,
message="Invalid Address {mac:{0}}".format(mac))
mc = self._machine_controls[mac]
try:
socket.inet_aton(mc['ip'])
except Exception as e:
return TrafficControlRc(
code=ReturnCode.INVALID_ADDRESS,
message="Invalid IP {0} for mac:{1}".format(mc['ip'], mac))
return method(self, mac, mc, *args)
def shapeMachine(self, mac, mc):
# remove old interface
if self._machine_shapings.get(mac):
# update shapings
old_shaping = self._machine_shapings[mac]
old_id = old_shaping['id']
self._unshape_interface(old_id, self.wan, old_shaping['ip'],
old_shaping['tc'].up)
self._unshape_interface(old_id, self.lan, old_shaping['ip'],
old_shaping['tc'].down)
self.idmanager.free(old_id)
del self._machine_shapings[mac]
# get profile setting
setting = self._profiles.get(mc.get('profile_name'))
if setting is None:
return TrafficControlRc(code=ACCESS_DENIED,
message="Invalid profile name: {0}".format(
mc['profile_name']))
new_id = None
try:
new_id = self.idmanager.new()
except Exception as e:
return TrafficControlRc(
code=ReturnCode.ID_EXHAUST,
message="No more session available: {0}".format(e))
self._machine_shapings[mac] = {
'id': new_id,
'ip': mc['ip'],
'tc': setting,
'profile_name': mc.get('profile_name')
}
# do shape
tcrc = self._shape_interface(new_id, self.wan, mc['ip'], setting.up)
if tcrc.code != ReturnCode.OK:
return tcrc
tcrc = self._shape_interface(new_id, self.lan, mc['ip'], setting.down)
if tcrc.code != ReturnCode.OK:
self._unshape_interface(new_id, self.wan, mc['ip'], setting.up)
return tcrc
mc['is_shaping'] = True
self._update_mcontrol(mc)
return TrafficControlRc(code=ReturnCode.OK)
def stopCapture(self, mac, mc):
self.logger.info("stopCapture mac:{0}".format(mac))
self.scapy_task.stopCapture(self.lan_name, mc['ip'], mac)
mc['is_capturing'] = False
self._update_mcontrol(mc)
return TrafficControlRc(code=ReturnCode.OK)
def startCapture(self, mac, mc, capture_filter):
self.logger.info("startCapture mac:{0} filter:{1}".format(
mac, capture_filter))
try:
self.scapy_task.startCapture(self.lan_name, capture_filter,
mc['ip'], mac)
except Exception as e:
self.logger.info("startCapture exception:{0}".format(str(e)))
mc['is_capturing'] = False
mc['capture_filter'] = None
self._update_mcontrol(mc)
return TrafficControlRc(code=ReturnCode.CAPTURE_NOT_READY,
message=str(e))
mc['is_capturing'] = True
mc['capture_filter'] = capture_filter
self._update_mcontrol(mc)
return TrafficControlRc(code=ReturnCode.OK)
def updateMachineControl(self, update_mc):
if not self._machine_controls.get(update_mc.mac):
# can not update while not exist
return TrafficControlRc(
code=ReturnCode.INVALID_ADDRESS,
message="Invalid Address {mac:{0}, ip:{1}}".format(
update_mc.mac, update_mc.state.ip))
mc = self._machine_controls[update_mc.mac]
# update profile_name only for now
mc['profile_name'] = update_mc.state.profile_name
self._update_mcontrol(mc)
# update profiles while shaping
if mc['is_shaping']:
return self.shapeMachine(update_mc.mac)
return TrafficControlRc(code=ReturnCode.OK)
def _shape_interface(self, mark, eth, ip, shaping):
"""Shape the traffic for a given interface.
Shape the traffic for a given IP on a given interface, given the mark
and the shaping settings.
There is a few steps to shape the traffic of an IP:
1. Create an HTB class that limit the throughput.
2. Create a NetEm QDisc that adds corruption, loss, reordering, loss
and delay.
3. Create the TC filter that will bucket packets with a given mark in
the right HTB class.
4. Set an iptables rule that mark packets going to/coming from IP
Args:
mark: The mark to set on IP packets.
eth: The network interface.
ip: The IP to shape traffic for.
shaping: The shaping setting to set.
Returns:
A TrafficControlRc containing information on success/failure.
"""
self.logger.info("Shaping ip {0} on interface {1}".format(
ip, eth['name']))
# HTB class
tcrc = self._set_htb_class(mark, eth, shaping)
if tcrc.code != ReturnCode.OK:
self.logger.error(
"adding HTB class on IFID {0}, mark {1}, err: {2}".format(
eth['name'], mark, tcrc.message))
return tcrc
# NetemQdisc
tcrc = self._set_netem_qdisc(mark, eth, shaping)
if tcrc.code != ReturnCode.OK:
self.logger.error(
"adding NetEm qdisc on IFID {0}, mark {1}, err: {2}".format(
eth['name'], mark, tcrc.message))
# delete class
self._unset_htb_class(mark, eth)
return tcrc
# filter
tcrc = self._set_filter(mark, eth, shaping)
if tcrc.code != ReturnCode.OK:
self.logger.error(
"adding filter FW on IFID {0}, mark {1}, err: {2}".format(
eth['name'], mark, tcrc.message))
# delete class
self._unset_htb_class(mark, eth)
return tcrc
# iptables
self._set_iptables(mark, eth, ip, shaping.iptables_options)
return TrafficControlRc(code=ReturnCode.OK)
def _set_htb_class(self, mark, eth, shaping):
"""Given a mark, an interface and shaping settings, set the HTB class.
Args:
mark: The mark based on which we create the class
eth: The interface on which to create that class id.
shaping: The shaping settings to set.
Returns:
A TrafficControlRc containing information on success/failure.
"""
ifid = eth['id']
idx = 0x10000 + mark
parent = 0x10000
self.logger.info("create new HTB class on IFID {0}, classid {1},"
"parent {2}, rate {3}kbits".format(
eth['name'], int_to_classid(idx),
int_to_classid(parent), shaping.rate
or 2**22 - 1))
try:
self.ipr.tc(
RTM_NEWTCLASS,
'htb',
ifid,
idx,
parent=parent,
rate="{}kbit".format(shaping.rate or (2**22 - 1)),
)
except NetlinkError as e:
return TrafficControlRc(code=ReturnCode.NETLINK_HTB_ERROR,
message=str(e))
except Exception as e:
self.logger.exception('_set_htb_class')
exc_info = sys.exc_info()
return TrafficControlRc(code=ReturnCode.UNKNOWN_HTB_ERROR,
message=str(exc_info))
return TrafficControlRc(code=ReturnCode.OK)
def unshapeMachine(self, mac):
# remove old interface
if self._machine_shapings.get(mac):
# update shapings
old_shaping = self._machine_shapings[mac]
old_id = old_shaping['id']
self._unshape_interface(old_id, self.wan, old_shaping['ip'],
old_shaping['tc'].up)
self._unshape_interface(old_id, self.lan, old_shaping['ip'],
old_shaping['tc'].down)
self.idmanager.free(old_id)
del self._machine_shapings[mac]
self._machine_controls[mac]['is_shaping'] = False
return TrafficControlRc(code=ReturnCode.OK)
def addProfile(self, profile):
self.logger.info('Request addProfile for name {0}'.format(
profile.name))
self.db_task.queue.put(
((profile.name, profile.tc_setting), 'add_profile'))
self._profiles[profile.name] = profile.tc_setting
mac_list = []
for mac in self._machine_shapings:
mshaping = self._machine_shapings[mac]
if mshaping['profile_name'] == profile.name:
mac_list.append(mac)
for mac in mac_list:
self.shapeMachine(mac) # reshape
return TrafficControlRc(code=ReturnCode.OK)
def removeProfile(self, name):
self.logger.info('Request removeProfile for name {0}'.format(name))
self.db_task.queue.put(((name), 'remove_profile'))
del self._profiles[name]
mac_list = []
for mac in self._machine_shapings:
mshaping = self._machine_shapings[mac]
if mshaping['profile_name'] == name:
mac_list.append(mac)
for mac in mac_list:
self.unshapeMachine(mac) # unshape
for mac in self._machine_controls:
mc = self._machine_controls[mac]
if mc.get('profile_name') and mc['profile_name'] == name:
del mc['profile_name']
self._update_mcontrol(mc)
return TrafficControlRc(code=ReturnCode.OK)
def _unshape_interface(self, mark, eth, ip, settings):
"""Unshape the traffic for a given interface.
Unshape the traffic for a given IP on a given interface, given the mark
and the shaping settings.
There is a few steps to unshape the traffic of an IP:
1. Remove the iptables rule.
2. Remove the TC filter.
3. Remove the HTB class.
Args:
mark: The mark to set on IP packets.
eth: The network interface.
ip: The IP to shape traffic for.
shaping: The shaping setting to set.
Returns:
A TrafficControlRc containing information on success/failure.
"""
self.logger.info("Unshaping ip {0} on interface {1}".format(
ip, eth['name']))
# iptables
self._unset_iptables(mark, eth, ip, settings.iptables_options)
# filter
tcrc = self._unset_filter(mark, eth)
if tcrc.code != ReturnCode.OK:
self.logger.error(
"deleting FW filter on IFID {0}, mark {1}, err: {2}".format(
eth['name'], mark, tcrc.message))
return tcrc
# HTB class
tcrc = self._unset_htb_class(mark, eth)
if tcrc.code != ReturnCode.OK:
self.logger.error(
"deleting HTB class on IFID {0}, mark {1}, err: {2}".format(
eth['name'], mark, tcrc.message))
return tcrc
return TrafficControlRc(code=ReturnCode.OK)
def _initialize_tc_for_interface(self, eth):
"""Initialize TC on a given interface.
If an exception is thrown, it will be forwarded to the main loop
unless it can be ignored.
Args:
eth: the interface to flush TC on.
Raises:
NetlinkError: An error occured initializing TC subsystem.
Exception: Any other exception thrown during initialization.
"""
idx = 0x10000
eth_name = eth['name']
eth_id = eth['id']
try:
self.logger.info("deleting root QDisc on {0}".format(eth_name))
self.ipr.tc(RTM_DELQDISC, None, eth_id, 0, parent=TC_H_ROOT)
except Exception as e:
# a (2, 'No such file or directory') can be thrown if there is
# nothing to delete. Ignore such error, return the error otherwise
if isinstance(e, NetlinkError) and e.code == 2:
self.logger.warning("could not delete root QDisc. There might "
"have been nothing to delete")
else:
self.logger.exception(
'Initializing root Qdisc for {0}'.format(eth_name))
raise
try:
self.logger.info("setting root qdisc on {0}".format(eth_name))
self.ipr.tc(RTM_NEWQDISC, "htb", eth_id, idx, default=0)
except Exception as e:
self.logger.exception(
'Setting root Qdisc for {0}'.format(eth_name))
raise
return TrafficControlRc(code=ReturnCode.OK)