Skip to content

Commit 5f66caa

Browse files
bdracobdraco
authored
Mark DNSOutgoing write functions as protected (#633)
1 parent 4ce33e4 commit 5f66caa

1 file changed

Lines changed: 40 additions & 44 deletions

File tree

‎zeroconf/_dns.py‎

Lines changed: 40 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -572,7 +572,7 @@ def __init__(self, flags: int, multicast: bool = True, id_: int = 0) -> None:
572572
self.multicast = multicast
573573
self.packets_data: List[bytes] = []
574574

575-
# these 3 are per-packet -- see also reset_for_next_packet()
575+
# these 3 are per-packet -- see also _reset_for_next_packet()
576576
self.names: Dict[str, int] = {}
577577
self.data: List[bytes] = []
578578
self.size: int = 12
@@ -585,7 +585,7 @@ def __init__(self, flags: int, multicast: bool = True, id_: int = 0) -> None:
585585
self.authorities: List[DNSPointer] = []
586586
self.additionals: List[DNSRecord] = []
587587

588-
def reset_for_next_packet(self) -> None:
588+
def _reset_for_next_packet(self) -> None:
589589
self.names = {}
590590
self.data = []
591591
self.size = 12
@@ -686,51 +686,51 @@ def add_question_or_all_cache(
686686
for cached_entry in cached_entries:
687687
self.add_answer_at_time(cached_entry, now)
688688

689-
def pack(self, format_: Union[bytes, str], value: Any) -> None:
689+
def _pack(self, format_: Union[bytes, str], value: Any) -> None:
690690
self.data.append(struct.pack(format_, value))
691691
self.size += struct.calcsize(format_)
692692

693-
def write_byte(self, value: int) -> None:
693+
def _write_byte(self, value: int) -> None:
694694
"""Writes a single byte to the packet"""
695-
self.pack(b'!c', int2byte(value))
695+
self._pack(b'!c', int2byte(value))
696696

697-
def insert_short_at_start(self, value: int) -> None:
697+
def _insert_short_at_start(self, value: int) -> None:
698698
"""Inserts an unsigned short at the start of the packet"""
699699
self.data.insert(0, struct.pack(b'!H', value))
700700

701-
def replace_short(self, index: int, value: int) -> None:
701+
def _replace_short(self, index: int, value: int) -> None:
702702
"""Replaces an unsigned short in a certain position in the packet"""
703703
self.data[index] = struct.pack(b'!H', value)
704704

705705
def write_short(self, value: int) -> None:
706706
"""Writes an unsigned short to the packet"""
707-
self.pack(b'!H', value)
707+
self._pack(b'!H', value)
708708

709-
def write_int(self, value: Union[float, int]) -> None:
709+
def _write_int(self, value: Union[float, int]) -> None:
710710
"""Writes an unsigned integer to the packet"""
711-
self.pack(b'!I', int(value))
711+
self._pack(b'!I', int(value))
712712

713713
def write_string(self, value: bytes) -> None:
714714
"""Writes a string to the packet"""
715715
assert isinstance(value, bytes)
716716
self.data.append(value)
717717
self.size += len(value)
718718

719-
def write_utf(self, s: str) -> None:
719+
def _write_utf(self, s: str) -> None:
720720
"""Writes a UTF-8 string of a given length to the packet"""
721721
utfstr = s.encode('utf-8')
722722
length = len(utfstr)
723723
if length > 64:
724724
raise NamePartTooLongException
725-
self.write_byte(length)
725+
self._write_byte(length)
726726
self.write_string(utfstr)
727727

728728
def write_character_string(self, value: bytes) -> None:
729729
assert isinstance(value, bytes)
730730
length = len(value)
731731
if length > 256:
732732
raise NamePartTooLongException
733-
self.write_byte(length)
733+
self._write_byte(length)
734734
self.write_string(value)
735735

736736
def write_name(self, name: str) -> None:
@@ -768,49 +768,45 @@ def write_name(self, name: str) -> None:
768768

769769
# write the new names out.
770770
for part in parts[:count]:
771-
self.write_utf(part)
771+
self._write_utf(part)
772772

773773
# if we wrote part of the name, create a pointer to the rest
774774
if count != len(name_suffices):
775775
# Found substring in packet, create pointer
776776
index = self.names[name_suffices[count]]
777-
self.write_byte((index >> 8) | 0xC0)
778-
self.write_byte(index & 0xFF)
777+
self._write_byte((index >> 8) | 0xC0)
778+
self._write_byte(index & 0xFF)
779779
else:
780780
# this is the end of a name
781-
self.write_byte(0)
781+
self._write_byte(0)
782782

783-
def write_question(self, question: DNSQuestion) -> bool:
783+
def _write_question(self, question: DNSQuestion) -> bool:
784784
"""Writes a question to the packet"""
785785
start_data_length, start_size = len(self.data), self.size
786786
self.write_name(question.name)
787787
self.write_short(question.type)
788-
self.write_record_class(question)
788+
self._write_record_class(question)
789789
return self._check_data_limit_or_rollback(start_data_length, start_size)
790790

791-
def write_record_class(self, record: Union[DNSQuestion, DNSRecord]) -> None:
791+
def _write_record_class(self, record: Union[DNSQuestion, DNSRecord]) -> None:
792792
"""Write out the record class including the unique/unicast (QU) bit."""
793793
if record.unique and self.multicast:
794794
self.write_short(record.class_ | _CLASS_UNIQUE)
795795
else:
796796
self.write_short(record.class_)
797797

798-
def write_record(self, record: DNSRecord, now: float) -> bool:
798+
def _write_record(self, record: DNSRecord, now: float) -> bool:
799799
"""Writes a record (answer, authoritative answer, additional) to
800-
the packet. Returns True on success, or False if we did not (either
801-
because the packet was already finished or because the record does
802-
not fit."""
803-
if self.state == self.State.finished:
804-
return False
805-
800+
the packet. Returns True on success, or False if we did not
801+
because the packet because the record does not fit."""
806802
start_data_length, start_size = len(self.data), self.size
807803
self.write_name(record.name)
808804
self.write_short(record.type)
809-
self.write_record_class(record)
805+
self._write_record_class(record)
810806
if now == 0:
811-
self.write_int(record.ttl)
807+
self._write_int(record.ttl)
812808
else:
813-
self.write_int(record.get_remaining_ttl(now))
809+
self._write_int(record.get_remaining_ttl(now))
814810
index = len(self.data)
815811

816812
self.write_short(0) # Will get replaced with the actual size
@@ -819,7 +815,7 @@ def write_record(self, record: DNSRecord, now: float) -> bool:
819815
length = sum((len(d) for d in self.data[index + 1 :]))
820816
# Here we replace the 0 length short we wrote
821817
# before with the actual length
822-
self.replace_short(index, length)
818+
self._replace_short(index, length)
823819
return self._check_data_limit_or_rollback(start_data_length, start_size)
824820

825821
def _check_data_limit_or_rollback(self, start_data_length: int, start_size: int) -> bool:
@@ -844,31 +840,31 @@ def _check_data_limit_or_rollback(self, start_data_length: int, start_size: int)
844840
def _write_questions_from_offset(self, questions_offset: int) -> int:
845841
questions_written = 0
846842
for question in self.questions[questions_offset:]:
847-
if not self.write_question(question):
843+
if not self._write_question(question):
848844
break
849845
questions_written += 1
850846
return questions_written
851847

852848
def _write_answers_from_offset(self, answer_offset: int) -> int:
853849
answers_written = 0
854850
for answer, time_ in self.answers[answer_offset:]:
855-
if not self.write_record(answer, time_):
851+
if not self._write_record(answer, time_):
856852
break
857853
answers_written += 1
858854
return answers_written
859855

860856
def _write_authorities_from_offset(self, authority_offset: int) -> int:
861857
authorities_written = 0
862858
for authority in self.authorities[authority_offset:]:
863-
if not self.write_record(authority, 0):
859+
if not self._write_record(authority, 0):
864860
break
865861
authorities_written += 1
866862
return authorities_written
867863

868864
def _write_additionals_from_offset(self, additional_offset: int) -> int:
869865
additionals_written = 0
870866
for additional in self.additionals[additional_offset:]:
871-
if not self.write_record(additional, 0):
867+
if not self._write_record(additional, 0):
872868
break
873869
additionals_written += 1
874870
return additionals_written
@@ -928,10 +924,10 @@ def packets(self) -> List[bytes]:
928924
authorities_written = self._write_authorities_from_offset(authority_offset)
929925
additionals_written = self._write_additionals_from_offset(additional_offset)
930926

931-
self.insert_short_at_start(additionals_written)
932-
self.insert_short_at_start(authorities_written)
933-
self.insert_short_at_start(answers_written)
934-
self.insert_short_at_start(questions_written)
927+
self._insert_short_at_start(additionals_written)
928+
self._insert_short_at_start(authorities_written)
929+
self._insert_short_at_start(answers_written)
930+
self._insert_short_at_start(questions_written)
935931

936932
questions_offset += questions_written
937933
answer_offset += answers_written
@@ -950,17 +946,17 @@ def packets(self) -> List[bytes]:
950946
):
951947
# https://datatracker.ietf.org/doc/html/rfc6762#section-7.2
952948
log.debug("Setting TC flag")
953-
self.insert_short_at_start(self.flags | _FLAGS_TC)
949+
self._insert_short_at_start(self.flags | _FLAGS_TC)
954950
else:
955-
self.insert_short_at_start(self.flags)
951+
self._insert_short_at_start(self.flags)
956952

957953
if self.multicast:
958-
self.insert_short_at_start(0)
954+
self._insert_short_at_start(0)
959955
else:
960-
self.insert_short_at_start(self.id)
956+
self._insert_short_at_start(self.id)
961957

962958
self.packets_data.append(b''.join(self.data))
963-
self.reset_for_next_packet()
959+
self._reset_for_next_packet()
964960

965961
if (questions_written + answers_written + authorities_written + additionals_written) == 0 and (
966962
len(self.questions) + len(self.answers) + len(self.authorities) + len(self.additionals)

0 commit comments

Comments
 (0)