ID2T-toolkit-BotnetTraffic/code/ID2TLib/MacAddressGenerator.py

from random import getrandbits

class MacAddressGenerator:
	def __init__(self, include_broadcast_macs = False, include_virtual_macs = False):
		self.broadcast = include_broadcast_macs
		self.virtual = include_virtual_macs
		
		self.generated = set()
	
	def random_mac(self) -> str:
		while True:
			mac = self._random_mac()
			if mac not in self.generated:
				self.generated.add(mac)
				return mac
	
	def clear(self):
		self.generated.clear()
	
	def generates_broadcast_macs(self) -> bool:
		return self.broadcast
	
	def generates_virtual_macs(self) -> bool:
		return self.virtual
	
	def set_broadcast_generation(self, broadcast: bool):
		self.broadcast = broadcast
	
	def set_virtual_generation(self, virtual: bool):
		self.virtual = virtual
	
	def _random_mac(self) -> str:
		mac_bytes = bytearray(getrandbits(8) for i in range(6))
		if not self.broadcast:
			mac_bytes[0] &= ~1 # clear the first bytes' first bit
		if not self.virtual:
			mac_bytes[0] &= ~2 # clear the first bytes' second bit
		
		return ":".join("%02X" % b for b in mac_bytes)