Merge pull request #19 from Ton-Dynasty/feat/dns-record

Add DNS resolution and address comparison functionality

examples/v3/get_dns_record.py

@@ -0,0 +1,22 @@
+import asyncio
+
+from pytoncenter import get_client
+from pytoncenter.v3.models import *
+
+
+async def main():
+    # DNS only supports mainnet, so you have to use mainnet api key here
+    # We set it to empty string to request without api key
+    client = get_client(version="v3", network="mainnet", api_key="")
+    dns_name = "doge.ton"
+    record = await client.get_dns_record(
+        GetDNSRecordRequest(
+            dns_name=dns_name,
+            category="wallet",
+        )
+    )
+    print(f"DNS record for {dns_name} is https://tonviewer.com/{record.address}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

pytoncenter/__init__.py

@@ -9,7 +9,20 @@ def get_client(version: Literal["v2"], network: Literal["mainnet", "testnet"], q
 
 
 @overload
-def get_client(version: Literal["v3"], network: Literal["mainnet", "testnet"], qps: Optional[float] = None, *args, **kwargs) -> AsyncTonCenterClientV3: ...
+def get_client(version: Literal["v3"], network: Literal["mainnet", "testnet"], qps: Optional[float] = None, *args, **kwargs) -> AsyncTonCenterClientV3:
+    """
+    Parameters
+    ----------
+    network : Union[Literal["mainnet"], Literal["testnet"]]
+        The network to use. Only mainnet and testnet are supported.
+
+    api_key : Optional[str], optional
+        The API key to use, by default None. If api_key is an empty string, then it will override the environment variable `TONCENTER_API_KEY`.
+    custom_endpoint : Optional[str], optional
+        The custom endpoint to use. If provided, it will override the network parameter.
+    qps: Optional[float], optional
+        The maximum queries per second to use. If not provided, it will use 9.5 if api_key is provided, otherwise 1.
+    """
 
 
 def get_client(version: Literal["v2", "v3"], network: Literal["mainnet", "testnet"], *args, **kwargs):

pytoncenter/address.py

@@ -228,9 +228,11 @@ def __key__(self) -> str:
         return self.to_string(True, True, True, is_test_only=False)
 
     def __eq__(self, __value: object) -> bool:
-        if not isinstance(__value, Address):
-            return False
-        return self.to_string(True, True, True, is_test_only=False) == __value.to_string(True, True, True, is_test_only=False)
+        if isinstance(__value, Address):
+            return self.to_string(True, True, True, is_test_only=False) == __value.to_string(True, True, True, is_test_only=False)
+        if isinstance(__value, str):
+            return self.to_string(True, True, True, is_test_only=False) == Address(__value).to_string(True, True, True, is_test_only=False)
+        return False
 
     def __repr__(self) -> str:
         return self.to_string(True, True, True, is_test_only=False)

pytoncenter/v3/api.py

@@ -1,10 +1,12 @@
 import asyncio
+import hashlib
 import os
 import time
 import warnings
 from typing import Any, Dict, List, Literal, Optional, Tuple, Union, overload
 
 import aiohttp
+from tonpy import CellSlice, begin_cell
 
 from pytoncenter.address import Address
 from pytoncenter.exception import TonCenterException, TonCenterValidationException
@@ -23,7 +25,25 @@ def __init__(
         qps: Optional[float] = None,
         **kwargs,
     ) -> None:
-        api_key = os.getenv("TONCENTER_API_KEY", api_key)
+        """
+        Parameters
+        ----------
+        network : Union[Literal["mainnet"], Literal["testnet"]]
+            The network to use. Only mainnet and testnet are supported.
+
+        api_key : Optional[str], optional
+            The API key to use, by default None. If api_key is an empty string, then it will override the environment variable `TONCENTER_API_KEY`.
+        custom_endpoint : Optional[str], optional
+            The custom endpoint to use. If provided, it will override the network parameter.
+        qps: Optional[float], optional
+            The maximum queries per second to use. If not provided, it will use 9.5 if api_key is provided, otherwise 1.
+        """
+        self._network = network
+        if api_key is not None:
+            assert isinstance(api_key, str), "API key must be a string"
+            self.api_key = api_key
+        else:
+            self.api_key = os.getenv("TONCENTER_API_KEY", None)
         # show warning if api_key is None
         if not api_key:
             warnings.warn(
@@ -36,7 +56,6 @@ def __init__(
         else:
             prefix = "" if network == "mainnet" else "testnet."
             self.base_url = f"https://{prefix}toncenter.com/api/v3"
-        self.api_key = api_key
 
         if qps is not None:
             assert qps > 0, "QPS must be greater than 0"
@@ -295,6 +314,48 @@ async def estimate_fee(self, req: EstimateFeeRequest) -> EstimateFeeResponse:
         resp = await self._async_post("estimateFee", req.model_dump(exclude_none=True))
         return EstimateFeeResponse(**resp)
 
+    async def get_dns_record(self, req: GetDNSRecordRequest) -> DNSRecord:
+        """
+        get_dns_record returns the DNS record for the given domain name and category.
+        Only mainnet is supported.
+        Only the wallet category is supported at the moment.
+        """
+        assert self._network == "mainnet", "Only mainnet is supported"
+        dns_name = req.dns_name.replace(".", "\0")
+        dns_slice = begin_cell().store_string(dns_name).end_cell().to_boc()
+        category_hash = 0
+        if req.category is not None:
+            category_hash = int(hashlib.sha256(req.category.encode()).hexdigest(), 16)
+        resp = await self.run_get_method(
+            RunGetMethodRequest(
+                address="EQC3dNlesgVD8YbAazcauIrXBPfiVhMMr5YYk2in0Mtsz0Bz",
+                method="dnsresolve",
+                stack=[
+                    GetMethodParameterInput(type="slice", value=dns_slice),
+                    GetMethodParameterInput(type="num", value=category_hash),
+                ],
+            )
+        )
+        # We cannot import decoder here because it will cause circular import
+        # Thus we need to handle the decoding here
+        # It might be dirty, will find a better way to handle this
+        assert len(resp.stack) == 2, "Expecting to find two items in the stack"
+        assert resp.stack[0].type == "num", "Expecting the first item to be a number"
+        assert resp.stack[1].type == "cell", "Expecting the second item to be a cell"
+        subdomain_bits = int(resp.stack[0].value, 16)  # type: ignore
+        cs = CellSlice(resp.stack[1].value)  # type: ignore
+        if req.category == "wallet":
+            _ = cs.load_uint(16)  # opcode
+            try:
+                address = cs.load_address()
+            except:
+                address = None
+            return DNSRecord(
+                subdomain_bits=subdomain_bits,
+                address=address,
+            )
+        raise NotImplementedError(f"Decoding {req.category} category is not implemented yet")
+
     async def wait_message_exists(self, req: WaitMessageExistsRequest):
         """
         wait_message_exists wait until the whole transaction trace is complete and yields the transaction.

pytoncenter/v3/models/customize.py

@@ -35,6 +35,8 @@
     "GetSourceTransactionRequest",
     "SubscribeTransactionRequest",
     "WaitMessageExistsRequest",
+    "GetDNSRecordRequest",
+    "DNSRecord",
 ]
 
 
@@ -315,3 +317,13 @@ class GetAccountRequest(BaseModel):
 
 class GetWalletRequest(BaseModel):
     address: AddressLike = Field(description="Account address. Account address. Can be sent in raw or user-friendly form")
+
+
+class GetDNSRecordRequest(BaseModel):
+    dns_name: str = Field(description="DNS name")
+    category: Optional[Literal["wallet"]] = Field(default="wallet", description="DNS category")
+
+
+class DNSRecord(BaseModel):
+    subdomain_bits: int = Field(description="Subdomain bits, for example doge.ton contains 32 bits (4 words)")
+    address: Optional[AddressLike] = Field(description="Address of the wallet")

tests/v2/address.py

@@ -1,4 +1,7 @@
+from typing import Union
+
 import pytest
+
 from pytoncenter.address import Address
 from pytoncenter.v2.api import AsyncTonCenterClientV2
 
@@ -22,6 +25,17 @@ def setup_method(self):
     def test_address_conversion(self, addr1: str, addr2: str, match: bool):
         assert (Address(addr1) == Address(addr2)) == match
 
+    @pytest.mark.parametrize(
+        ("addr1", "addr2", "match"),
+        [
+            (Address("0:2b790db779a6e344ee6094b09b859e0ac50f523888edc3678cd8fb845d784865"), "kQAreQ23eabjRO5glLCbhZ4KxQ9SOIjtw2eM2PuEXXhIZeh3", True),
+            ("kQAreQ23eabjRO5glLCbhZ4KxQ9SOIjtw2eM2PuEXXhIZeh3", Address("0QAreQ23eabjRO5glLCbhZ4KxQ9SOIjtw2eM2PuEXXhIZbWy"), True),
+            (Address("kQAreQ23eabjRO5glLCbhZ4KxQ9SOIjtw2eM2PuEXXhIZeh3"), "EQCpk40ub48fvx89vSUjOTRy0vOEEZ4crOPPfLEvg88q1PwN", False),
+        ],
+    )
+    def test_address_compare_eq(self, addr1: Union[str, Address], addr2: Union[str, Address], match: bool):
+        assert (addr1 == addr2) == match
+
     @pytest.mark.parametrize(
         ("addr", "form"),
         [

tests/v3/decode.py

@@ -1,8 +1,11 @@
+import asyncio
+
 import pytest
-from pytoncenter import get_client, AsyncTonCenterClientV3
+
+from pytoncenter import AsyncTonCenterClientV3, get_client
 from pytoncenter.address import Address
-from pytoncenter.decoder import JettonDataDecoder, Types, Decoder, AutoDecoder
-from pytoncenter.v3.models import RunGetMethodRequest
+from pytoncenter.decoder import AutoDecoder, Decoder, JettonDataDecoder, Types
+from pytoncenter.v3.models import GetDNSRecordRequest, RunGetMethodRequest
 
 pytest_plugins = ("pytest_asyncio",)
 
@@ -55,3 +58,16 @@ async def test_auto_decoder(self):
         assert output["idx_2"] == 9
         assert output["idx_7"] == -1  # True, because auto decoder does not know the type of the field
         assert output["idx_4"] == 1000000000
+
+    @pytest.mark.asyncio
+    async def test_get_dns_record(self):
+        await asyncio.sleep(0.5)
+        client = get_client(version="v3", network="mainnet", api_key="")
+        dns_name = "doge.ton"
+        record = await client.get_dns_record(
+            GetDNSRecordRequest(
+                dns_name=dns_name,
+                category="wallet",
+            )
+        )
+        assert record.address == Address("EQDVjQWmoS6xrPqPJ5vEFBPZdBnY075ydcoEEqpVWjJXZ9RE")