Add wait logic

SRT/netfunnel.py

@@ -11,9 +11,13 @@ class NetFunnelHelper:
 
     OP_CODE = {
         "getTidchkEnter": "5101",
+        "chkEnter": "5002",
         "setComplete": "5004",
     }
 
+    WAIT_STATUS_PASS = "200"  # No need to wait
+    WAIT_STATUS_FAIL = "201"  # Need to wait
+
     DEFAULT_HEADERS = {
         "User-Agent": USER_AGENT,
         "Accept": "*/*",
@@ -30,7 +34,7 @@ class NetFunnelHelper:
     def __init__(self):
         self.session = requests.session()
         self.session.headers.update(self.DEFAULT_HEADERS)
-        self._cachedNetfunnelKey = None
+        self._cached_key = None
 
     def generate_netfunnel_key(self, use_cache: bool):
         key = self._get_netfunnel_key(use_cache)
@@ -48,8 +52,8 @@ def _get_netfunnel_key(self, use_cache: bool):
             str: NetFunnel 키
         """
 
-        if use_cache and self._cachedNetfunnelKey is not None:
-            return self._cachedNetfunnelKey
+        if use_cache and self._cached_key is not None:
+            return self._cached_key
 
         params = {
             "opcode": self.OP_CODE["getTidchkEnter"],
@@ -62,18 +66,77 @@ def _get_netfunnel_key(self, use_cache: bool):
         }
 
         try:
-            response = self.session.get(
+            resp = self.session.get(
                 self.NETFUNNEL_URL,
                 params=params,
-            ).text
+            )
         except Exception as e:
             raise SRTNetFunnelError(e) from e
 
-        netfunnel_key = self._extract_netfunnel_key(response)
-        self._cachedNetfunnelKey = netfunnel_key
+        netfunnel_resp = NetFunnelResponse.parse(resp.text)
+
+        netfunnel_key = netfunnel_resp.get("key")
+        if netfunnel_key is None:
+            raise SRTNetFunnelError("NetFunnel key not found in response")
+
+        if netfunnel_resp.get("status") == self.WAIT_STATUS_FAIL:
+            # TODO: better logging
+            print("접속자가 많아 대기열에 들어갑니다.")
+
+            nwait = netfunnel_resp.get("nwait") or "<unknown>"
+
+            netfunnel_key = self._wait_until_complete(netfunnel_key, nwait)
+
+        self._cached_key = netfunnel_key
 
         return netfunnel_key
 
+    def _wait_until_complete(self, key: str, nwait: str) -> str:
+        """
+        NetFunnel이 완료될 때까지 대기합니다.
+        """
+
+        params = {
+            "opcode": self.OP_CODE["chkEnter"],
+            "key": key,
+            "nfid": "0",
+            "prefix": f"NetFunnel.gRtype={self.OP_CODE['chkEnter']};",
+            "ttl": 1,
+            "sid": "service_1",
+            "aid": "act_10",
+            "js": "true",
+            self._get_timestamp_for_netfunnel(): "",
+        }
+
+        try:
+            resp = self.session.get(
+                self.NETFUNNEL_URL,
+                params=params,
+            )
+        except Exception as e:
+            raise SRTNetFunnelError(e) from e
+
+        netfunnel_resp = NetFunnelResponse.parse(resp.text)
+
+        if netfunnel_resp.get("status") != self.WAIT_STATUS_PASS:
+            key = netfunnel_resp.get("key")
+            if key is None:
+                raise SRTNetFunnelError("NetFunnel key not found in response")
+
+            nwait = netfunnel_resp.get("nwait") or "<unknown>"
+
+
+            print(f"대기인원: {nwait}명")
+
+            # 1 sec
+            # TODO: find how to calculate the re-try interval
+            time.sleep(1)
+
+            return self._wait_until_complete(key, nwait)
+        else:
+            return key
+
+
     def _set_complete(self, key: str):
         """
         NetFunnel 완료 요청을 보냅니다.
@@ -92,36 +155,89 @@ def _set_complete(self, key: str):
         }
 
         try:
-            self.session.get(
+            resp = self.session.get(
                 self.NETFUNNEL_URL,
                 params=params,
             )
+
+            netfunnel_resp = NetFunnelResponse.parse(resp.text)
+            if netfunnel_resp.get("status") != self.WAIT_STATUS_PASS:
+                raise SRTNetFunnelError(f"Failed to complete NetFunnel: {netfunnel_resp}")
+
         except Exception as e:
             raise SRTNetFunnelError(e) from e
 
     def _get_timestamp_for_netfunnel(self):
         return int(time.time() * 1000)
 
-    def _extract_netfunnel_key(self, response: str):
+class NetFunnelResponse:
+    """
+    Represents a NetFunnel response.
+    """
+
+    OP_CODE_KEY = "NetFunnel.gRtype"
+    RESULT_KEY = "NetFunnel.gControl.result"
+
+    RESULT_SUBKEYS = [
+        "key",
+        "nwait",
+        "nnext",
+        "tps",
+        "ttl",
+        "ip",
+        "port",
+        "msg",
+    ]
+
+    def __init__(self, response: str, data: dict[str, str]):
+        self.response = response
+        self.data = data
+
+    @classmethod
+    def parse(cls, response: str) -> "NetFunnelResponse":
         """
-        NetFunnel 키를 추출합니다.
-
-        Args:
-            response (str): NetFunnel 응답
+        The factory method to create a NetFunnelResponse object from a response string.
 
-        Returns:
-            str: NetFunnel 키
+        the response string will be in the format of:
 
-        Raises:
-            SRTNetFunnelError: NetFunnel 키를 찾을 수 없는 경우
+        ```
+        NetFunnel.gRtype=5101;
+        NetFunnel.gControl.result='5002:201:key=<key>&nwait=3&nnext=1&tps=11.247706&ttl=1&ip=<ip>&port=80';
+        NetFunnel.gControl._showResult();
+        ```
         """
-        key_start = response.find("key=")
-        if key_start == -1:
-            raise SRTNetFunnelError("NetFunnel key not found in response")
 
-        key_start += 4  # "key=" length
-        key_end = response.find("&", key_start)
-        if key_end == -1:
-            raise SRTNetFunnelError("Invalid NetFunnel key format in response")
+        top_level_keys = [r.strip() for r in response.split(";")]
+
+        data: dict[str, str] = {}
+        for top_level_key in top_level_keys:
+
+            if top_level_key.startswith(cls.OP_CODE_KEY):
+                data["opcode"] = top_level_key[len(cls.OP_CODE_KEY) + 1 :]
+
+            if top_level_key.startswith(cls.RESULT_KEY):
+                results = top_level_key[len(cls.RESULT_KEY) + 1 :].strip("'").split(":")
+
+                if len(results) != 3:
+                    raise SRTNetFunnelError(f"Invalid NetFunnel response format: {response}")
+
+                code, status, result = results
+                data["next_code"] = code  # dunno what this is...
+                data["status"] = status
+
+                result_keys = result.split("&")
+                for key in result_keys:
+                    for subkey in cls.RESULT_SUBKEYS:
+                        if key.startswith(subkey):
+                            data[subkey] = key.split("=")[1]
+                            break
+
+        return cls(response, data)
+
+    def get(self, key: str) -> None:
+        return self.data.get(key)
+
+    def __str__(self):
+        return self.response
+
 
-        return response[key_start:key_end]

tests/test_netfunnel.py

@@ -1,6 +1,6 @@
 import pytest
 
-from SRT.netfunnel import NetFunnelHelper, SRTNetFunnelError
+from SRT.netfunnel import NetFunnelHelper, SRTNetFunnelError, NetFunnelResponse
 
 
 @pytest.fixture(scope="module")
@@ -25,12 +25,52 @@ def test_set_complete_success(helper):
     assert True
 
 
-def test_extract_netfunnel_key_success(helper):
+def test_netfunnel_response_parse():
     key = "C75890BD44561ED79DFA180832F2D016F95C9F7BE965B2"
-    response = f"""
-NetFunnel.gRtype = 5101;
-NetFunnel.gControl.result = '5002:200:key={key}&nwait=0&nnext=0&tps=0&ttl=0&ip=nf.letskorail.com&port=443';
+    response_text = f"""
+NetFunnel.gRtype=5101;
+NetFunnel.gControl.result='5002:200:key={key}&nwait=0&nnext=0&tps=0&ttl=0&ip=nf.letskorail.com&port=443';
 NetFunnel.gControl._showResult();
 """
-    key = helper._extract_netfunnel_key(response)
-    assert key == key
+    response = NetFunnelResponse.parse(response_text)
+    assert "key" in response.data
+    assert response.data["key"] == key
+    assert response.data["nwait"] == "0"
+    assert response.data["nnext"] == "0"
+    assert response.data["tps"] == "0"
+    assert response.data["ttl"] == "0"
+    assert response.data["ip"] == "nf.letskorail.com"
+    assert response.data["port"] == "443"
+    assert response.data["opcode"] == "5101"
+    assert response.data["next_code"] == "5002"
+    assert response.data["status"] == "200"
+
+    response_text = """
+NetFunnel.gRtype=5002;
+NetFunnel.gControl.result='5002:200:key=test_key&nwait=0&nnext=0&tps=0&ttl=0&ip=nf.letskorail.com&port=443';
+NetFunnel.gControl._showResult();
+"""
+
+    response = NetFunnelResponse.parse(response_text)
+    assert "key" in response.data
+    assert response.data["key"] == "test_key"
+    assert response.data["nwait"] == "0"
+    assert response.data["nnext"] == "0"
+    assert response.data["tps"] == "0"
+    assert response.data["ttl"] == "0"
+    assert response.data["ip"] == "nf.letskorail.com"
+    assert response.data["port"] == "443"
+    assert response.data["opcode"] == "5002"
+    assert response.data["next_code"] == "5002"
+    assert response.data["status"] == "200"
+
+    response_text = """
+NetFunnel.gRtype=5004;NetFunnel.gControl.result='5004:502:msg="Already Completed"'; NetFunnel.gControl._showResult();
+"""
+
+    response = NetFunnelResponse.parse(response_text)
+    assert "key" not in response.data
+    assert response.data["opcode"] == "5004"
+    assert response.data["next_code"] == "5004"
+    assert response.data["status"] == "502"
+    assert response.data["msg"] == '"Already Completed"'