Refactor http module

Use an API client; cleaner code; get rid of PATHS from constants.

Author: jleclanche

bna/constants.py

@@ -9,10 +9,3 @@
 	# "US": "us.mobile-service.blizzard.com",
 	"default": "mobile-service.blizzard.com",
 }
-
-PATHS = {
-	"init_restore": "/enrollment/initiatePaperRestore.htm",
-	"validate_restore": "/enrollment/validatePaperRestore.htm",
-	"enroll": "/enrollment/enroll.htm",
-	"time": "/enrollment/time.htm",
-}

bna/http.py

@@ -7,7 +7,7 @@
 from time import time
 from typing import Optional, Tuple
 
-from .constants import ENROLL_HOSTS, PATHS
+from .constants import ENROLL_HOSTS
 from .crypto import decrypt, encrypt, restore_code_to_bytes
 from .utils import normalize_serial
 
@@ -18,27 +18,48 @@ def __init__(self, msg, response):
 		super().__init__(msg)
 
 
-def _post(host: str, path: str, *, data: Optional[str] = None) -> bytes:
-	"""
-	Send computed data to Blizzard servers
-	Return the answer from the server
-	"""
-	conn = HTTPConnection(host)
-	conn.request("POST", path, data)
-	response = conn.getresponse()
+class APIClient:
+	def __init__(self, *, region: str = "US", host: str = ""):
+		self.region = region
+		self.host = host or ENROLL_HOSTS.get(region, ENROLL_HOSTS["default"])
+
+	def post(self, path: str, *, data: Optional[str] = None) -> bytes:
+		conn = HTTPConnection(self.host)
+		conn.request("POST", path, data)
+		response = conn.getresponse()
+
+		if response.status != 200:
+			raise HTTPError("%s returned status %i" % (self.host, response.status), response)
 
-	if response.status != 200:
-		raise HTTPError("%s returned status %i" % (host, response.status), response)
+		ret = response.read()
+		conn.close()
+		return ret
 
-	ret = response.read()
-	conn.close()
-	return ret
+	def enroll(self, data):
+		return self.post("/enrollment/enroll.htm", data=data)
 
+	def get_time(self) -> int:
+		response = self.post("/enrollment/time.htm")
+		return int(struct.unpack(">Q", response)[0])
 
-def enroll(
-	data: str, host: str = ENROLL_HOSTS["default"], path: str = PATHS["enroll"]
-) -> bytes:
-	return _post(host, path, data=data)
+	def initiate_paper_restore(self, serial: str):
+		response = self.post("/enrollment/initiatePaperRestore.htm", data=serial)
+		resp_size = len(response)
+		if resp_size != 32:
+			raise ValueError("Bad challenge response (%i bytes)" % (resp_size))
+
+		return response
+
+	def validate_paper_restore(self, serial: str, encrypted_data: str):
+		data = serial + encrypted_data
+		try:
+			response = self.post("/enrollment/validatePaperRestore.htm", data=data)
+		except HTTPError as e:
+			if e.response.status == 600:
+				raise HTTPError("Invalid serial or restore key", e.response)
+			else:
+				raise
+		return response
 
 
 def request_new_serial(
@@ -57,14 +78,15 @@ def base_msg(otp, region, model):
 
 	otp = token_bytes(37)
 	data = base_msg(otp, region, model)
+	encrypted_data = encrypt(data)
+
+	client = APIClient(region=region)
+	response = client.enroll(encrypted_data)[8:]
 
-	e = encrypt(data)
-	# get the host, or fallback to default
-	host = ENROLL_HOSTS.get(region, ENROLL_HOSTS["default"])
-	response = decrypt(enroll(e, host)[8:], otp)
+	decrypted_response = decrypt(response, otp)
 
-	secret = b32encode(response[:20]).decode()
-	serial = response[20:].decode()
+	secret = b32encode(decrypted_response[:20]).decode()
+	serial = decrypted_response[20:].decode()
 
 	region = serial[:2]
 	if region not in ("CN", "EU", "US"):
@@ -73,7 +95,7 @@ def base_msg(otp, region, model):
 	return serial, secret
 
 
-def get_time_offset(region: str = "US", path: str = PATHS["time"]) -> int:
+def get_time_offset(region: str = "US") -> int:
 	"""
 	Calculates the time difference in seconds as a float
 	between the local host and a remote server
@@ -82,52 +104,30 @@ def get_time_offset(region: str = "US", path: str = PATHS["time"]) -> int:
 	Negative numbers indicate the local clock is ahead of the
 	server clock.
 	"""
-	host = ENROLL_HOSTS.get(region, ENROLL_HOSTS["default"])
-	response = _post(host, path)
-	t = time()
+	client = APIClient(region=region)
+	server_time = client.get_time()
+	local_time = time()
 
 	# NOTE: The server returns time in milliseconds as an int whereas
 	# Python returns it as a float, in seconds.
-	server_time = int(struct.unpack(">Q", response)[0])
-
-	return server_time - int(t * 1000)
+	return server_time - int(local_time * 1000)
 
 
 def restore(serial: str, restore_code: str) -> str:
-	restore_code = restore_code.upper()
 	serial = normalize_serial(serial)
+	restore_code = restore_code.upper()
 	if len(restore_code) != 10:
 		raise ValueError(f"invalid restore code (should be 10 characters): {restore_code}")
 
-	challenge = initiate_paper_restore(serial)
-	if len(challenge) != 32:
-		raise ValueError("Bad challenge length (expected 32, got %i)" % (len(challenge)))
+	client = APIClient()
+	challenge = client.initiate_paper_restore(serial)
 
 	code = restore_code_to_bytes(restore_code)
 	hash = hmac.new(code, serial.encode() + challenge, digestmod=sha1).digest()
 
 	otp = token_bytes(20)
-	e = encrypt(hash + otp)
-	response = validate_paper_restore(serial + e)
+	encrypted_data = encrypt(hash + otp)
+	response = client.validate_paper_restore(serial, encrypted_data)
 	secret = decrypt(response, otp)
 
 	return b32encode(secret).decode()
-
-
-def initiate_paper_restore(
-	serial: str, host: str = ENROLL_HOSTS["default"], path: str = PATHS["init_restore"]
-) -> bytes:
-	return _post(host, path, data=serial)
-
-
-def validate_paper_restore(
-	data: str, host: str = ENROLL_HOSTS["default"], path: str = PATHS["validate_restore"]
-) -> bytes:
-	try:
-		response = _post(host, path, data=data)
-	except HTTPError as e:
-		if e.response.status == 600:
-			raise HTTPError("Invalid serial or restore key", e.response)
-		else:
-			raise
-	return response