Tighten up some logic/flow in the new ETM packet code

1 files changed, 12 insertions, 23 deletions

diff --git a/paramiko/packet.py b/paramiko/packet.py
index 6395ce6c..12663168 100644
--- a/paramiko/packet.py
+++ b/paramiko/packet.py
@@ -420,14 +420,8 @@ class Packetizer(object):
                 out = packet
             # + mac
             if self.__block_engine_out is not None:
-                if self.__etm_out:
-                    payload = (
-                        struct.pack(">I", self.__sequence_number_out) + out
-                    )
-                else:
-                    payload = (
-                        struct.pack(">I", self.__sequence_number_out) + packet
-                    )
+                packed = struct.pack(">I", self.__sequence_number_out)
+                payload = packed + (out if self.__etm_out else packet)
                 out += compute_hmac(
                     self.__mac_key_out, payload, self.__mac_engine_out
                 )[: self.__mac_size_out]
@@ -484,17 +478,16 @@ class Packetizer(object):
         if self.__dump_packets:
             self._log(DEBUG, util.format_binary(header, "IN: "))
 
-        # already computed
-        packet_size = (
-            packet_size
-            if self.__etm_in
-            else struct.unpack(">I", header[:4])[0]
-        )
-        # leftover contains decrypted bytes from the first block (after the
-        # length field)
-
-        # no leftovers
-        if not self.__etm_in:
+        # When ETM is in play, we've already read the packet size & decrypted
+        # everything, so just set the packet back to the header we obtained.
+        if self.__etm_in:
+            packet = header
+        # Otherwise, use the older non-ETM logic
+        else:
+            packet_size = struct.unpack(">I", header[:4])[0]
+
+            # leftover contains decrypted bytes from the first block (after the
+            # length field)
             leftover = header[4:]
             if (packet_size - len(leftover)) % self.__block_size_in != 0:
                 raise SSHException("Invalid packet blocking")
@@ -508,10 +501,6 @@ class Packetizer(object):
                 packet = self.__block_engine_in.update(packet)
             packet = leftover + packet
 
-        else:
-            # already decrypted everything above
-            packet = header
-
         if self.__dump_packets:
             self._log(DEBUG, util.format_binary(packet, "IN: "))