Skip to content

Commit

Permalink
Refactor ICMPv6 option handling to use Icmp6NdOptions class.
Browse files Browse the repository at this point in the history
  • Loading branch information
@ccie18643
ccie18643 committed Jul 15, 2024
1 parent ca8f187 commit c0e73a4
Show file tree
Hide file tree
Showing 14 changed files with 353 additions and 323 deletions.
83 changes: 7 additions & 76 deletions pytcp/protocols/icmp6/assembler.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@

# pylint: disable = too-many-instance-attributes
# pylint: disable = too-many-locals
# pylint: disable = too-few-public-methods
# pylint: disable = too-many-return-statements
# pylint: disable = too-many-arguments
# pylint: disable = redefined-builtin
Expand All @@ -41,9 +42,8 @@
from __future__ import annotations

import struct
from typing import TYPE_CHECKING

from pytcp.lib.ip6_address import Ip6Address, Ip6Network
from pytcp.lib.ip6_address import Ip6Address
from pytcp.lib.ip_helper import inet_cksum
from pytcp.lib.proto import ProtoAssembler
from pytcp.lib.tracker import Tracker
Expand Down Expand Up @@ -79,17 +79,9 @@
from pytcp.protocols.icmp6.message__nd__router_solicitation import (
Icmp6NdRouterSolicitationMessage,
)
from pytcp.protocols.icmp6.options__nd import (
Icmp6NdOption,
Icmp6NdOptionPi,
Icmp6NdOptionSlla,
Icmp6NdOptionTlla,
)
from pytcp.protocols.icmp6.options__nd import Icmp6NdOptions
from pytcp.protocols.ip6.header import IP6_HEADER_LEN

if TYPE_CHECKING:
from pytcp.lib.mac_address import MacAddress


class Icmp6Assembler(Icmp6, ProtoAssembler):
"""
Expand Down Expand Up @@ -216,7 +208,7 @@ class Icmp6NdRouterSolicitationMessageAssembler(
Assembler class for the ICMPv6 ND Router Soliciation message.
"""

def __init__(self, *, nd_options: list[Icmp6NdOption]) -> None:
def __init__(self, *, nd_options: Icmp6NdOptions) -> None:
"""
Create the message object.
"""
Expand All @@ -241,7 +233,7 @@ def __init__(
router_lifetime: int,
reachable_time: int,
retrans_timer: int,
nd_options: list[Icmp6NdOption],
nd_options: Icmp6NdOptions,
) -> None:
"""
Create the message object.
Expand Down Expand Up @@ -269,7 +261,7 @@ class Icmp6NdNeighborSolicitationMessageAssembler(
"""

def __init__(
self, *, target_address: Ip6Address, nd_options: list[Icmp6NdOption]
self, *, target_address: Ip6Address, nd_options: Icmp6NdOptions
) -> None:
"""
Create the message object.
Expand All @@ -294,7 +286,7 @@ def __init__(
flag_s: bool = False,
flag_o: bool = False,
target_address: Ip6Address,
nd_options: list[Icmp6NdOption],
nd_options: Icmp6NdOptions,
) -> None:
"""
Create the message object.
Expand Down Expand Up @@ -323,67 +315,6 @@ def __init__(self, *, records: list[Icmp6Mld2AddressRecord]) -> None:
self._records = records


#
# ICMPv6 Neighbor Discovery options
#


class Icmp6NdOptSllaAssembler(Icmp6NdOptionSlla):
"""
ICMPv6 ND option assembler - Source Link Layer Address (1).
"""

def __init__(self, *, slla: MacAddress) -> None:
"""
Create the option object.
"""

self._slla = slla


class Icmp6NdOptTllaAssembler(Icmp6NdOptionTlla):
"""
ICMPv6 ND option - Target Link Layer Address (2).
"""

def __init__(self, *, tlla: MacAddress) -> None:
"""
Create the option object.
"""

self._tlla = tlla


class Icmp6NdOptPiAssembler(Icmp6NdOptionPi):
"""
ICMPv6 ND option assembler - Prefix Information (3).
"""

def __init__(
self,
*,
flag_l: bool = False,
flag_a: bool = False,
flag_r: bool = False,
valid_lifetime: int,
prefer_lifetime: int,
prefix: Ip6Network,
) -> None:
"""
Option constructor.
"""

assert 0 <= valid_lifetime <= 0xFFFFFFFF
assert 0 <= prefer_lifetime <= 0xFFFFFFFF

self._flag_l = flag_l
self._flag_a = flag_a
self._flag_r = flag_r
self._valid_lifetime = valid_lifetime
self._prefer_lifetime = prefer_lifetime
self._prefix = prefix


#
# ICMPv6 MLD2 Multicast support classes
#
Expand Down
53 changes: 8 additions & 45 deletions pytcp/protocols/icmp6/message__nd.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,15 +38,9 @@ class for the ICMPv6 protocol.

from __future__ import annotations

from pytcp.lib.ip6_address import Ip6Network
from pytcp.lib.mac_address import MacAddress
from pytcp.protocols.icmp6.message import Icmp6Message, Icmp6Type
from pytcp.protocols.icmp6.options__nd import (
Icmp6NdOption,
Icmp6NdOptionPi,
Icmp6NdOptionSlla,
Icmp6NdOptionTlla,
)
from pytcp.protocols.icmp6.options__nd import Icmp6NdOptions, NdPrefixInfo


class Icmp6NdMessage(Icmp6Message):
Expand All @@ -55,20 +49,12 @@ class Icmp6NdMessage(Icmp6Message):
"""

_type: Icmp6Type
_nd_options: list[Icmp6NdOption]
_nd_options: Icmp6NdOptions

@property
def _raw_nd_options(self) -> bytes:
def nd_options(self) -> Icmp6NdOptions:
"""
Get the ND options in raw format.
"""

return b"".join(bytes(nd_option) for nd_option in self._nd_options)

@property
def nd_options(self) -> list[Icmp6NdOption]:
"""
Get the '_nd_options' property.
Get the '_nd_options' attribute.
"""

return self._nd_options
Expand All @@ -79,43 +65,20 @@ def option_slla(self) -> MacAddress | None:
Get the Source Link Layer Address option.
"""

# TODO: Currently we are returning the first option,
# should we return list of matching options instead ?

for option in self._nd_options:
if isinstance(option, Icmp6NdOptionSlla):
return option.slla

return None
return self._nd_options.slla

@property
def option_tlla(self) -> MacAddress | None:
"""
Get the Target Link Layer Address option.
"""

# TODO: Currently we are returning the first option,
# should we return list of matching options instead ?

for option in self._nd_options:
if isinstance(option, Icmp6NdOptionTlla):
return option.tlla

return None
return self._nd_options.tlla

@property
def option_pi(self) -> list[Ip6Network]:
def option_pi(self) -> list[NdPrefixInfo]:
"""
Get the Prefix Information option.
"""

# TODO: Currently we are returning the first option,
# should we return list of matching options instead ?

prefixes = []

for option in self._nd_options:
if isinstance(option, Icmp6NdOptionPi):
prefixes.append(option.prefix)

return prefixes
return self._nd_options.pi
16 changes: 8 additions & 8 deletions pytcp/protocols/icmp6/message__nd__neighbor_advertisement.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ class for the ICMPv6 protocol.
from pytcp.lib.ip6_address import Ip6Address
from pytcp.protocols.icmp6.message import Icmp6Code, Icmp6Type
from pytcp.protocols.icmp6.message__nd import Icmp6NdMessage
from pytcp.protocols.icmp6.options__nd import Icmp6NdOption
from pytcp.protocols.icmp6.options__nd import Icmp6NdOptions

# 'ND - Neighbor Advertisement' message (136/0)

Expand Down Expand Up @@ -86,16 +86,16 @@ class Icmp6NdNeighborAdvertisementMessage(Icmp6NdMessage):
_flag_s: bool
_flag_o: bool
_target_address: Ip6Address
_nd_options: list[Icmp6NdOption]
_nd_options: Icmp6NdOptions

@override
def __len__(self) -> int:
"""
Get the message length.
"""

return ICMP6_MESSAGE_LEN__ND_NEIGHBOR_ADVERTISEMENT + sum(
len(option) for option in self._nd_options
return ICMP6_MESSAGE_LEN__ND_NEIGHBOR_ADVERTISEMENT + len(
self._nd_options
)

@override
Expand All @@ -111,7 +111,7 @@ def __str__(self) -> str:
f"{'S' if self._flag_s else '-'}"
f"{'O' if self._flag_o else '-'}, "
f"target {self._target_address}, "
f"{', '.join(str(nd_option) for nd_option in self._nd_options)}"
f"opts [{self._nd_options}]"
)

@override
Expand All @@ -126,7 +126,7 @@ def __repr__(self) -> str:
f"flag_s={self._flag_s!r}, "
f"flag_o={self._flag_o!r}, "
f"target={self._target_address!r}, "
f"{', '.join(f"{nd_option!r}" for nd_option in self._nd_options)})"
f"options={self._nd_options!r})"
)

@override
Expand All @@ -136,7 +136,7 @@ def __bytes__(self) -> bytes:
"""

return struct.pack(
f"! BBH L 16s {len(self._raw_nd_options)}s",
f"! BBH L 16s {len(self._nd_options)}s",
int(self._type),
int(self._code),
0,
Expand All @@ -145,7 +145,7 @@ def __bytes__(self) -> bytes:
| (self._flag_o << 29)
| (self._reserved & 0b00011111_11111111_11111111_11111111),
bytes(self._target_address),
self._raw_nd_options,
bytes(self._nd_options),
)

@property
Expand Down
18 changes: 9 additions & 9 deletions pytcp/protocols/icmp6/message__nd__neighbor_solicitation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ class for the ICMPv6 protocol.
from pytcp.lib.ip6_address import Ip6Address
from pytcp.protocols.icmp6.message import Icmp6Code, Icmp6Type
from pytcp.protocols.icmp6.message__nd import Icmp6NdMessage
from pytcp.protocols.icmp6.options__nd import Icmp6NdOption
from pytcp.protocols.icmp6.options__nd import Icmp6NdOptions

# 'Neighbor Discovery - Neighbor Solicitation' message (135/0)

Expand Down Expand Up @@ -83,16 +83,16 @@ class Icmp6NdNeighborSolicitationMessage(Icmp6NdMessage):

_reserved: int
_target_address: Ip6Address
_nd_options: list[Icmp6NdOption]
_nd_options: Icmp6NdOptions

@override
def __len__(self) -> int:
"""
Get the message length.
"""

return ICMP6_MESSAGE_LEN__ND_NEIGHBOR_SOLICITATION + sum(
len(option) for option in self._nd_options
return ICMP6_MESSAGE_LEN__ND_NEIGHBOR_SOLICITATION + len(
self._nd_options
)

@override
Expand All @@ -104,7 +104,7 @@ def __str__(self) -> str:
return (
f"ICMP6 ND Neighbor Solicitation, "
f"target {self._target_address}, "
f"{', '.join(str(nd_option) for nd_option in self._nd_options)}"
f"opts [{self._nd_options}]"
)

@override
Expand All @@ -115,8 +115,8 @@ def __repr__(self) -> str:

return (
"Icmp6NdNeighborSolicitationMessage("
f"target {self._target_address!r}, "
f"{', '.join(f"{nd_option!r}" for nd_option in self._nd_options)})"
f"target={self._target_address!r}, "
f"options={self._nd_options!r})"
)

@override
Expand All @@ -126,13 +126,13 @@ def __bytes__(self) -> bytes:
"""

return struct.pack(
f"! BBH L 16s {len(self._raw_nd_options)}s",
f"! BBH L 16s {len(self._nd_options)}s",
int(self._type),
int(self._code),
0,
self._reserved,
bytes(self._target_address),
self._raw_nd_options,
bytes(self._nd_options),
)

@property
Expand Down
Loading

0 comments on commit c0e73a4

Please sign in to comment.