From 77e947f77b4887274a4b165efc966456b9caaee6 Mon Sep 17 00:00:00 2001
From: Markus Birth <mbirth@gmail.com>
Date: Sat, 10 Feb 2018 03:30:10 +0100
Subject: [PATCH] Added new style EncryptHeaderRequest. tclcheck.py completely
 working again.

---
 tclcheck.py                             | 30 +++++++++++++---------
 tcllib/requests/__init__.py             |  1 +
 tcllib/requests/encryptheaderrequest.py | 34 +++++++++++++++++++++++++
 tcllib/requests/tclresult.py            |  4 +++
 4 files changed, 57 insertions(+), 12 deletions(-)
 create mode 100644 tcllib/requests/encryptheaderrequest.py

diff --git a/tclcheck.py b/tclcheck.py
index 6ae2eeb..2ec366a 100755
--- a/tclcheck.py
+++ b/tclcheck.py
@@ -9,10 +9,11 @@ import os
 import random
 import sys
 
-import tcllib
-import tcllib.argparser
+from tcllib import argparser
 from tcllib.devices import Device
-from tcllib.requests import RequestRunner, CheckRequest, DownloadRequest, ChecksumRequest, ServerSelector, write_info_if_dumps_found
+from tcllib.requests import RequestRunner, CheckRequest, DownloadRequest, \
+        ChecksumRequest, EncryptHeaderRequest, ServerSelector, \
+        write_info_if_dumps_found
 from tcllib.xmltools import pretty_xml
 
 
@@ -20,7 +21,7 @@ dpdesc = """
     Checks for the latest FULL updates for the specified PRD number or for an OTA from the
     version specified as fvver.
     """
-dp = tcllib.argparser.DefaultParser(__file__, dpdesc)
+dp = argparser.DefaultParser(__file__, dpdesc)
 dp.add_argument("prd", nargs=1, help="CU Reference #, e.g. PRD-63117-011")
 dp.add_argument("fvver", nargs="?", help="Firmware version to check for OTA updates, e.g. AAM481 (omit to run FULL check)", default="AAA000")
 dp.add_argument("-i", "--imei", help="use specified IMEI instead of default", type=str)
@@ -88,9 +89,9 @@ dlrres = dlr.get_result()
 print(dlrres.pretty_xml())
 
 if dlrres.encslaves:
-    cksrunner = RequestRunner(ServerSelector(dlrres.encslaves), https=False)
+    encrunner = RequestRunner(ServerSelector(dlrres.encslaves), https=False)
     cks = ChecksumRequest(dlrres.fileurl, dlrres.fileurl)
-    cksrunner.run(cks)
+    encrunner.run(cks)
     if not cks.success:
         print("{}".format(cks.error))
         sys.exit(4)
@@ -104,17 +105,22 @@ for s in dlrres.s3_slaves:
     print("http://{}{}".format(s, dlrres.s3_fileurl))
 
 if dev.mode == dev.MODE_STATES["FULL"]:
-    header = fc.do_encrypt_header(random.choice(encslaves), fileurl)
-    headname = "header_{}.bin".format(tv)
+    hdr = EncryptHeaderRequest(dlrres.fileurl)
+    encrunner.run(hdr)
+    if not hdr.success:
+        print("{}".format(hdr.error))
+        sys.exit(5)
+    hdrres = hdr.get_result()
+    headname = "header_{}.bin".format(chkres.tvver)
     headdir = "headers"
     if not os.path.exists(headdir):
         os.makedirs(headdir)
-    if len(header) == 4194320:
+    if len(hdrres.rawdata) == 4194320:
         # TODO: Check sha1sum
         print("Header length check passed. Writing to {}.".format(headname))
         with open(os.path.join(headdir, headname), "wb") as f:
-            f.write(header)
+            f.write(hdrres.rawdata)
     else:
-        print("Header length invalid ({}).".format(len(header)))
+        print("Header length invalid ({}).".format(len(hdrres.rawdata)))
 
-tcllib.FotaCheck.write_info_if_dumps_found()
+write_info_if_dumps_found()
diff --git a/tcllib/requests/__init__.py b/tcllib/requests/__init__.py
index 93e7cd3..7d4dc86 100644
--- a/tcllib/requests/__init__.py
+++ b/tcllib/requests/__init__.py
@@ -3,6 +3,7 @@
 from .checkrequest import CheckRequest
 from .downloadrequest import DownloadRequest
 from .checksumrequest import ChecksumRequest
+from .encryptheaderrequest import EncryptHeaderRequest
 from .runner import *
 from .serverselector import *
 from .dumpmgr import write_info_if_dumps_found
diff --git a/tcllib/requests/encryptheaderrequest.py b/tcllib/requests/encryptheaderrequest.py
new file mode 100644
index 0000000..e78f06a
--- /dev/null
+++ b/tcllib/requests/encryptheaderrequest.py
@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+
+from collections import OrderedDict
+from .. import credentials, devices
+from .tclrequest import TclRequest
+from .tclresult import EncryptHeaderResult
+
+class EncryptHeaderRequest(TclRequest):
+    def __init__(self, file_uri):
+        super().__init__()
+        # NOTE: THIS HAS TO BE RUN ON AN ENCSLAVE
+        self.uri = "/encrypt_header.php"
+        self.rawmode = True
+        self.method = "POST"
+        self.file_uri = file_uri
+
+    def get_headers(self):
+        return {"User-Agent": "tcl"}
+
+    def get_params(self):
+        params = OrderedDict()
+        params.update(credentials.get_creds2())
+        params["address"] = bytes(self.file_uri, "utf-8")
+        return params
+
+    def is_done(self, http_status: int, contents: str) -> bool:
+        # Expect "HTTP 206 Partial Content" response
+        if http_status == 206:
+            self.result = EncryptHeaderResult(contents)
+            self.success = True
+            return True
+        self.error = "HTTP {}".format(http_status)
+        self.success = False
+        return True
diff --git a/tcllib/requests/tclresult.py b/tcllib/requests/tclresult.py
index 146fba0..335a44d 100644
--- a/tcllib/requests/tclresult.py
+++ b/tcllib/requests/tclresult.py
@@ -62,3 +62,7 @@ class ChecksumResult(TclResult):
         self.sha1_enc_footer = file.find("ENCRYPT_FOOTER").text
         self.sha1_footer = file.find("FOOTER").text
         self.sha1_body = file.find("BODY").text
+
+class EncryptHeaderResult(TclResult):
+    def __init__(self, contents: str):
+        self.rawdata = contents