diff --git a/cloudbaseinit/tests/utils/test_dhcp.py b/cloudbaseinit/tests/utils/test_dhcp.py index 9298cee7..5604d672 100644 --- a/cloudbaseinit/tests/utils/test_dhcp.py +++ b/cloudbaseinit/tests/utils/test_dhcp.py @@ -12,6 +12,7 @@ # License for the specific language governing permissions and limitations # under the License. +import binascii import mock import netifaces import struct @@ -28,9 +29,9 @@ class DHCPUtilsTests(unittest.TestCase): def test_get_dhcp_request_data(self): - fake_mac_address = '01:02:03:04:05:06' - fake_mac_address_b = bytearray( - fake_mac_address.replace(':', '').decode('hex')) + fake_mac_address = '010203040506' + fake_mac_address_a = fake_mac_address.encode('ascii', 'strict') + fake_mac_address_b = bytearray(binascii.unhexlify(fake_mac_address_a)) data = b'\x01' data += b'\x01' @@ -66,12 +67,12 @@ class DHCPUtilsTests(unittest.TestCase): @mock.patch('struct.unpack') def _test_parse_dhcp_reply(self, mock_unpack, message_type, id_reply, equals_cookie): - fake_data = 236 * "1" + fake_data = 236 * b"1" if equals_cookie: - fake_data += dhcp._DHCP_COOKIE + '11' + fake_data += dhcp._DHCP_COOKIE + b'11' else: - fake_data += '111111' - fake_data += 'fake' + fake_data += b'111111' + fake_data += b'fake' fake_data += dhcp._OPTION_END mock_unpack.side_effect = [(message_type, None), (id_reply, None), @@ -86,7 +87,7 @@ class DHCPUtilsTests(unittest.TestCase): elif fake_data[236:240] != dhcp._DHCP_COOKIE: self.assertEqual(response, (False, {})) else: - self.assertEqual(response, (True, {100: 'fake'})) + self.assertEqual(response, (True, {100: b'fake'})) def test_parse_dhcp_reply(self): self._test_parse_dhcp_reply(message_type=2, id_reply=9999, diff --git a/cloudbaseinit/utils/dhcp.py b/cloudbaseinit/utils/dhcp.py index 22efe28e..7feb0eff 100644 --- a/cloudbaseinit/utils/dhcp.py +++ b/cloudbaseinit/utils/dhcp.py @@ -12,6 +12,7 @@ # License for the specific language governing permissions and limitations # under the License. +import binascii import datetime import netifaces import random @@ -33,7 +34,8 @@ LOG = logging.getLogger(__name__) def _get_dhcp_request_data(id_req, mac_address, requested_options, vendor_id): - mac_address_b = bytearray(mac_address.replace(':', '').decode('hex')) + mac_address_ascii = mac_address.replace(':', '').encode('ascii', 'strict') + mac_address_b = bytearray(binascii.unhexlify(mac_address_ascii)) # See: http://www.ietf.org/rfc/rfc2131.txt data = b'\x01' @@ -69,7 +71,7 @@ def _get_dhcp_request_data(id_req, mac_address, requested_options, def _parse_dhcp_reply(data, id_req): - message_type = struct.unpack('b', data[0])[0] + message_type = struct.unpack('b', data[0:1])[0] if message_type != 2: return (False, {}) @@ -84,9 +86,10 @@ def _parse_dhcp_reply(data, id_req): options = {} i = 240 - while data[i] != _OPTION_END: - id_option = struct.unpack('b', data[i])[0] - option_data_len = struct.unpack('b', data[i + 1])[0] + data_len = len(data) + while i < data_len and data[i:i + 1] != _OPTION_END: + id_option = struct.unpack('b', data[i:i + 1])[0] + option_data_len = struct.unpack('b', data[i + 1:i + 2])[0] i += 2 options[id_option] = data[i: i + option_data_len] i += option_data_len