From 79703c96ef4141115111241b859e058667088ca0 Mon Sep 17 00:00:00 2001
From: Dimitri Papadopoulos
 <3350651-DimitriPapadopoulos@users.noreply.gitlab.com>
Date: Wed, 22 Sep 2021 22:27:50 +0200
Subject: [PATCH] Avoid code duplication

This fixes a CodeClimate alert.

Signed-off-by: Dimitri Papadopoulos <3350651-DimitriPapadopoulos@users.noreply.gitlab.com>
---
 trojans/tncc-emulate.py | 55 +++++++++++++++++------------------------
 1 file changed, 22 insertions(+), 33 deletions(-)

diff --git a/trojans/tncc-emulate.py b/trojans/tncc-emulate.py
index d80abe54..dd443e0d 100755
--- a/trojans/tncc-emulate.py
+++ b/trojans/tncc-emulate.py
@@ -94,9 +94,7 @@ MSG_FUNK_PLATFORM = 0x58301
 MSG_FUNK = 0xa4c01
 
 
-# 0013 - Message
-def decode_0013(buf, indent):
-    logging.debug('%scmd 0013 (Message) %d bytes', indent, len(buf))
+def _decode_helper(buf, indent):
     ret = collections.defaultdict(list)
     while (len(buf) >= 12):
         length, cmd, out = decode_packet(buf, indent + "  ")
@@ -104,6 +102,11 @@ def decode_0013(buf, indent):
         ret[cmd].append(out)
     return ret
 
+# 0013 - Message
+def decode_0013(buf, indent):
+    logging.debug('%scmd 0013 (Message) %d bytes', indent, len(buf))
+    return _decode_helper(buf, indent)
+
 
 # 0012 - u32
 def decode_0012(buf, indent):
@@ -116,23 +119,13 @@ def decode_0016(buf, indent):
     logging.debug('%scmd 0016 (compressed message) %d bytes', indent, len(buf))
     _, compressed = struct.unpack(">I" + str(len(buf) - 4) + "s", buf)
     buf = zlib.decompress(compressed)
-    ret = collections.defaultdict(list)
-    while (len(buf) >= 12):
-        length, cmd, out = decode_packet(buf, indent + "  ")
-        buf = buf[length:]
-        ret[cmd].append(out)
-    return ret
+    return _decode_helper(buf, indent)
 
 
 # 0ce4 - encapsulation
 def decode_0ce4(buf, indent):
     logging.debug('%scmd 0ce4 (encapsulation) %d bytes', indent, len(buf))
-    ret = collections.defaultdict(list)
-    while (len(buf) >= 12):
-        length, cmd, out = decode_packet(buf, indent + "  ")
-        buf = buf[length:]
-        ret[cmd].append(out)
-    return ret
+    return _decode_helper(buf, indent)
 
 
 # 0ce5 - string without hex prefixer
@@ -193,24 +186,20 @@ def decode_packet(buf, indent=""):
     if length % 4:
         length += 4 - (length % 4)
 
-    if cmd == 0x0013:
-        data = decode_0013(data, indent)
-    elif cmd == 0x0012:
-        data = decode_0012(data, indent)
-    elif cmd == 0x0016:
-        data = decode_0016(data, indent)
-    elif cmd == 0x0ce4:
-        data = decode_0ce4(data, indent)
-    elif cmd == 0x0ce5:
-        data = decode_0ce5(data, indent)
-    elif cmd == 0x0ce7:
-        data = decode_0ce7(data, indent)
-    elif cmd == 0x0cf0:
-        data = decode_0cf0(data, indent)
-    elif cmd == 0x0cf1:
-        data = decode_0cf1(data, indent)
-    elif cmd == 0x0cf3:
-        data = decode_0cf3(data, indent)
+    decode_function = {
+        0x0012: decode_0012,
+        0x0013: decode_0013,
+        0x0016: decode_0016,
+        0x0ce4: decode_0ce4,
+        0x0ce5: decode_0ce5,
+        0x0ce7: decode_0ce7,
+        0x0cf0: decode_0cf0,
+        0x0cf1: decode_0cf1,
+        0x0cf3: decode_0cf3,
+    }
+
+    if cmd in decode_function:
+        data = decode_function[cmd](data, indent)
     else:
         logging.debug('%scmd %04x(%02x:%02x) is unknown, length %d', indent, cmd, _1, _2, length)
         data = None
-- 
2.49.0