Improve login handling

mu5001tool/__main__.py

@@ -61,7 +61,7 @@ def main():
         m.set_password(args.password)
 
     if args.password_file is not None:
-        password = args.password_file.read_text()
+        password = args.password_file.read_text().strip()
         m.set_password(password)
 
     if args.password is not None or args.password_file is not None:

mu5001tool/api.py

@@ -3,9 +3,22 @@ import hashlib
 import requests
 import urllib
 
+class Mu5001ToolException(Exception):
+    pass
+
+class FailedToLogin(Mu5001ToolException):
+    pass
+
+class InvalidPassword(FailedToLogin):
+    pass
+
+class AccountLocked(FailedToLogin):
+    pass
+
 class Mu5001Tool:
     def __init__(self, host="http://192.168.0.1"):
         self._password = None
+        self._password_invalid = False
 
         self.host = host
         self.session = requests.Session()
@@ -49,6 +62,54 @@ class Mu5001Tool:
 
         return login_sha256(login_sha256(password) + ld)
 
+    def is_login_possible(self):
+        data = self.get_cmd_process({
+            "multi_data": "1",
+            "cmd": "login_lock_time,psw_fail_num_str,loginfo",
+        })
+
+        if data.get("loginfo") == "ok":
+            # Already logged in
+            return False
+
+        # How many password attempts left before locking
+        psw_fail_num_str = data.get("psw_fail_num_str")
+
+        if psw_fail_num_str == "":
+            psw_fail_num_str = None
+        else:
+            psw_fail_num_str = int(psw_fail_num_str)
+
+        if psw_fail_num_str > 0:
+            return True
+
+        # How long login is blocked
+        login_lock_time = data.get("login_lock_time")
+
+        if login_lock_time == "":
+            login_lock_time = None
+        else:
+            login_lock_time = int(login_lock_time)
+
+        # Login is blocked
+        return False
+
+
+    def login_if_possible(self):
+        if self._password is None:
+            return
+
+        if self._password_invalid:
+            return
+
+        if not self.is_login_possible():
+            return
+
+        try:
+            self.login()
+        except InvalidPassword as e:
+            self._password_invalid = True
+
     def login(self, password=None):
         if password is None:
             password = self._password
@@ -56,11 +117,22 @@ class Mu5001Tool:
         if password is None:
             raise ValueError("No password provided")
 
-        return self.set_cmd_process({
+        data = self.set_cmd_process({
             "goformId": "LOGIN",
             "password": self.hashed_login_password(password),
         })
 
+        result = data.get("result")
+
+        if result == "3":
+            raise InvalidPassword("Invalid password")
+
+        if result == "1":
+            raise AccountLocked("Account currently locked")
+
+        if result != "0":
+            raise FailedToLogin("Failed to log in")
+
     def logoff(self):
         return self.set_cmd_process({"goformId": "LOGOFF"})
 

mu5001tool/prometheus_exporter.py

@@ -1,5 +1,6 @@
 from . import Mu5001Tool
 from http.server import BaseHTTPRequestHandler, HTTPServer, HTTPStatus
+from pprint import pprint
 import socket
 import traceback
 
@@ -88,6 +89,7 @@ def make_prometheus_exporter_request_handler(m):
                 "battery_charg_type",
                 "data_volume_limit_unit",
                 "dial_mode",
+                "loginfo",
                 "mode_main_state",
                 "modem_main_state",
                 "network_provider",
@@ -116,6 +118,9 @@ def make_prometheus_exporter_request_handler(m):
                 "cmd": ",".join(cmds_as_metric_value + cmds_as_metric_label),
             })
 
+            if data.get("loginfo") != "ok":
+                m.login_if_possible()
+
             out = []
 
             for cmd in cmds_as_metric_value: