Add Support for Custom HTTP Headers in agent-python via Argument

__main__.py

@@ -15,7 +15,7 @@
 from htpclient.generic_cracker import GenericCracker
 from htpclient.hashcat_cracker import HashcatCracker
 from htpclient.hashlist import Hashlist
-from htpclient.helpers import start_uftpd, file_get_contents
+from htpclient.helpers import start_uftpd, file_get_contents, parse_http_headers
 from htpclient.initialize import Initialize
 from htpclient.jsonRequest import *
 from htpclient.dicts import *
@@ -140,6 +140,8 @@ def init(args):
         CONFIG.set_value('zaps-path', os.path.abspath(args.zaps_path))
     if args.preprocessors_path and len(args.preprocessors_path):
         CONFIG.set_value('preprocessors-path', os.path.abspath(args.preprocessors_path))
+    if args.http_headers and len(args.http_headers):
+        CONFIG.set_value('http-headers', parse_http_headers(args.http_headers))
 
     logging.info("Starting client '" + Initialize.get_version() + "'...")
 
@@ -153,6 +155,7 @@ def init(args):
 
     session = Session(requests.Session()).s
     session.headers.update({'User-Agent': Initialize.get_version()})
+    session.headers.update(CONFIG.get_value('http-headers'))
 
     if CONFIG.get_value('proxies'):
         session.proxies = CONFIG.get_value('proxies')
@@ -334,6 +337,7 @@ def de_register():
     parser.add_argument('--preprocessors-path', type=str, required=False, help='Use given folder path as preprocessors location')
     parser.add_argument('--zaps-path', type=str, required=False, help='Use given folder path as zaps location')
     parser.add_argument('--cpu-only', action='store_true', help='Force client to register as CPU only and also only reading out CPU information')
+    parser.add_argument('--http-headers', type=str, required=False, help='Use given http headers in all HTTP requests')
     args = parser.parse_args()
 
     if args.version:

build.sh

@@ -4,14 +4,16 @@ if [ -f hashtopolis.zip ]; then
   rm hashtopolis.zip
 fi
 
-# write commit count since release into version number when compiling into zip
-count=$(git log $(git describe --tags --abbrev=0)..HEAD --oneline | wc -l)
-if [ ${count} \> 0 ];
-then
-    sed -i -E 's/return "([0-9]+)\.([0-9]+)\.([0-9]+)"/return "\1.\2.\3.'$count'"/g' htpclient/initialize.py
-fi;
+# Get latest tag, fallback to initial commit hash if no tags found
+latest_tag=$(git describe --tags --abbrev=0 2>/dev/null || echo $(git rev-list --max-parents=0 HEAD))
+count=$(git log ${latest_tag}..HEAD --oneline | wc -l)
+
+if [ "$count" -gt 0 ]; then
+  sed -i -E 's/return "([0-9]+)\.([0-9]+)\.([0-9]+)"/return "\1.\2.\3.'$count'"/g' htpclient/initialize.py
+fi
+
 zip -r hashtopolis.zip __main__.py htpclient -x "*__pycache__*"
-if [ ${count} \> 0 ];
-then
-    sed -i -E 's/return "([0-9]+)\.([0-9]+)\.([0-9]+)\.([0-9]+)"/return "\1.\2.\3"/g' htpclient/initialize.py
-fi;
+
+if [ "$count" -gt 0 ]; then
+  sed -i -E 's/return "([0-9]+)\.([0-9]+)\.([0-9]+)\.([0-9]+)"/return "\1.\2.\3"/g' htpclient/initialize.py
+fi

htpclient/helpers.py

@@ -63,7 +63,8 @@ def file_get_contents(filename):
 
 def start_uftpd(os_extension, config):
     try:
-        subprocess.check_output("killall -s 9 uftpd", shell=True)  # stop running service to make sure we can start it again
+        # stop running service to make sure we can start it again
+        subprocess.check_output("killall -s 9 uftpd", shell=True)
     except subprocess.CalledProcessError:
         pass  # ignore in case uftpd was not running
     path = './uftpd' + os_extension
@@ -132,3 +133,14 @@ def update_files(command, prince=False):
 def escape_ansi(line):
     ansi_escape = re.compile(r'(\x9B|\x1B\[)[0-?]*[ -/]*[@-~]')
     return ansi_escape.sub('', line)
+
+
+def parse_http_headers(header_str):
+    headers = {}
+    if header_str:
+        pairs = header_str.split(',')
+        for pair in pairs:
+            if ':' in pair:
+                key, value = pair.split(':', 1)
+                headers[key.strip()] = value.strip()
+    return headers

htpclient/jsonRequest.py

@@ -13,7 +13,8 @@ def __init__(self, data):
     def execute(self):
         try:
             logging.debug(self.data)
-            r = self.session.post(self.config.get_value('url'), json=self.data, timeout=30)
+            headers = self.config.get_value('http-headers') or {}
+            r = self.session.post(self.config.get_value('url'), json=self.data, timeout=30, headers=headers)
             if r.status_code != 200:
                 logging.error("Status code from server: " + str(r.status_code))
                 return None