From 820375214f1b7b26109e5c3aea8d005fcc56eebf Mon Sep 17 00:00:00 2001
From: nlscc <66028747+nlscc@users.noreply.github.com>
Date: Tue, 1 Jun 2021 14:25:02 +0100
Subject: auto-normalize version codes with decrypt command (fix #48)

---
 samloader/crypt.py        |  4 +++-
 samloader/versionfetch.py | 17 ++++++++++-------
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/samloader/crypt.py b/samloader/crypt.py
index 0b57c5a..29611a6 100644
--- a/samloader/crypt.py
+++ b/samloader/crypt.py
@@ -8,8 +8,9 @@ import xml.etree.ElementTree as ET
 from Cryptodome.Cipher import AES
 from clint.textui import progress
 
-from . import request
 from . import fusclient
+from . import request
+from . import versionfetch
 
 # PKCS#7 unpad
 unpad = lambda d: d[:-d[-1]]
@@ -17,6 +18,7 @@ unpad = lambda d: d[:-d[-1]]
 def getv4key(version, model, region):
     """ Retrieve the AES key for V4 encryption. """
     client = fusclient.FUSClient()
+    version = versionfetch.normalizevercode(version)
     req = request.binaryinform(version, model, region, client.nonce)
     resp = client.makereq("NF_DownloadBinaryInform.do", req)
     root = ET.fromstring(resp)
diff --git a/samloader/versionfetch.py b/samloader/versionfetch.py
index b8f150c..e758867 100644
--- a/samloader/versionfetch.py
+++ b/samloader/versionfetch.py
@@ -6,6 +6,15 @@
 import xml.etree.ElementTree as ET
 import requests
 
+def normalizevercode(vercode: str) -> str:
+    """ Normalize a version code to four-part form. """
+    ver = vercode.split("/")
+    if len(ver) == 3:
+        ver.append(ver[0])
+    if ver[2] == "":
+        ver[2] = ver[0]
+    return "/".join(ver)
+
 def getlatestver(model: str, region: str) -> str:
     """ Get the latest firmware version code for a model and region. """
     req = requests.get("https://fota-cloud-dn.ospserver.net/firmware/" \
@@ -15,10 +24,4 @@ def getlatestver(model: str, region: str) -> str:
     vercode = root.find("./firmware/version/latest").text
     if vercode is None:
         raise Exception("No latest firmware found")
-    # Normalize retrieved version
-    ver = vercode.split("/")
-    if len(ver) == 3:
-        ver.append(ver[0])
-    if ver[2] == "":
-        ver[2] = ver[0]
-    return "/".join(ver)
+    return normalizevercode(vercode)
-- 
cgit v1.2.3