def request_dispatcher(file_path):
if "__" in file_path:
neighbor_host, neighbor_if = file_path.split('__')
neighbor_if = neighbor_if.split('.')[0]
for interface in devices[neighbor_host]['links']:
if neighbor_if == interface['interface_brief']:
ztp_host = interface['neighbor']
ztp_device = {}
ztp_device['hostname'] = ztp_host
ztp_device['staging_user'] = staging_user
ztp_device['staging_password'] = staging_password
ztp_device.update(devices[ztp_host])
config = render_file("base-configuration", **ztp_device)
return StringResponseData(config)
elif file_path == 'network-confg':
config = render_file(file_path,
staging_user=staging_user,
staging_password=staging_password)
return StringResponseData(config)
else:
return TftpData(file_path)
def testNextBlock(self):
class MockResponse:
def __init__(self, dataiter):
self._dataiter = dataiter
def read(self, size=0):
try:
return next(self._dataiter)
except StopIteration:
return None
# single-packet file
self.handler._last_block_sent = 0
self.handler._block_size = 1400
self.handler._response_data = StringResponseData("bacon")
self.handler._next_block()
self.assertEqual(self.handler._current_block, b"bacon")
self.assertEqual(self.handler._last_block_sent, 1)
# multi-packet file
self.handler._last_block_sent = 0
self.handler._block_size = 1400
self.handler._response_data = StringResponseData("bacon" * 281)
self.handler._next_block()
self.assertEqual(self.handler._current_block, b"bacon" * 280)
self.assertEqual(self.handler._last_block_sent, 1)
self.handler._next_block()
self.assertEqual(self.handler._current_block, b"bacon")
self.assertEqual(self.handler._last_block_sent, 2)
# partial read
data = MockResponse(iter("bacon"))
self.handler._last_block_sent = 0
self.handler._block_size = 1400
self.handler._response_data.read = data.read
self.handler._next_block()
self.assertEqual(self.handler._current_block, "bacon")
self.assertEqual(self.handler._last_block_sent, 1)
self.handler._last_block_sent = constants.MAX_BLOCK_NUMBER + 1
self.handler._next_block()
self.assertEqual(self.handler._last_block_sent, 0)
self.handler._response_data.read = Mock(side_effect=Exception("boom!"))
self.handler._next_block()
self.assertEqual(
self.handler._stats.error,
{
"error_code": constants.ERR_UNDEFINED,
"error_message": "Error while reading from source",
},
)
self.assertTrue(self.handler._should_stop)
def testParseOptionsNetascii(self):
self.handler._response_data = StringResponseData("foo\nbar\n")
self.handler._parse_options()
self.assertEqual(
self.handler._stats.options_in,
{
"mode": "netascii",
"blksize": 1492,
"tsize": 0,
"timeout": 99
},
)
self.assertIsInstance(self.handler._response_data, NetasciiReader)
self.assertEqual(self.handler._stats.blksize, 1492)
# options acked by the server don't include the mode
expected_opts_to_ack = self.options
del expected_opts_to_ack["mode"]
# tsize include the number of bytes in the response
expected_opts_to_ack["tsize"] = str(self.handler._response_data.size())
self.assertEqual(self.handler._stats.options, expected_opts_to_ack)
self.assertEqual(self.handler._stats.options_acked,
expected_opts_to_ack)
self.assertEqual(self.handler._tsize,
int(expected_opts_to_ack["tsize"]))
def request_dispatcher(file_path):
if file_path == 'network-confg':
config = render_file(file_path,
staging_user=staging_user,
staging_password=staging_password)
return StringResponseData(config)
else:
return TftpData(file_path)
def __init__(self,
server_addr,
peer,
path,
options,
stats_callback,
network_queue=()):
self.response = StringResponseData("foo")
super().__init__(server_addr, peer, path, options, stats_callback)
self.network_queue = network_queue
self.peer = peer
self._listener = MockSocketListener(network_queue, peer)
self._listener.sendto = Mock()
self._listener.close = Mock()
self._listener.settimeout = Mock()
def testNetAsciiReader(self):
tests = [
# content, expected output
("foo\nbar\nand another\none",
bytearray(b'foo\r\nbar\r\nand another\r\none')),
("foo\r\nbar\r\nand another\r\none",
bytearray(b'foo\r\x00\r\nbar\r\x00\r\nand another\r\x00\r\none')),
]
for input_content, expected in tests:
with self.subTest(content=input_content):
resp_data = StringResponseData(input_content)
n = NetasciiReader(resp_data)
self.assertGreater(n.size(), len(input_content))
output = n.read(512)
self.assertEqual(output, expected)
n.close()
def testNetAsciiReaderBig(self):
input_content = "I\nlike\ncrunchy\nbacon\n"
for _ in range(5):
input_content += input_content
resp_data = StringResponseData(input_content)
n = NetasciiReader(resp_data)
self.assertGreater(n.size(), 0)
self.assertGreater(n.size(), len(input_content))
block_size = 512
output = bytearray()
while True:
c = n.read(block_size)
output += c
if len(c) < block_size:
break
self.assertEqual(input_content.count("\n"), output.count(b"\r\n"))
n.close()
def get_response_data(self):
capath = self._capath
target = self._url
path = self._path
if self._path[0] == '/':
path = self._path[1:]
# accept only mac based config and defaults file
if path.startswith("pxelinux.cfg/01-"):
# request specific system configuration
logging.debug(f"Got request for {path}, forwarding to HTTP")
return HttpResponseData(f"https://{target}/tftp/{path}", capath)
elif path.startswith("pxelinux.cfg/default"):
# server local default
logging.debug(f"Got request for {path}, forwarding to HTTP")
return HttpResponseData(f"https://{target}/tftp/{path}", capath)
elif path.startswith("pxelinux.cfg/"):
# ignore other pxelinux.cfg files
logging.debug(f"Got request for {path}, ignoring")
return StringResponseData("")
# The rest get from http
logging.debug(f"Got request for {path}, forwarding to HTTP")
return HttpResponseData(f"https://{target}/tftp/{path}", capath)