Skip to content

Commit bdd32b4

Browse files
authored
use dns resolver to get account key (#12)
* use dns resolver to get account key * add comment for get_key() * remove default value for first_mapping_account_key in PythClient * remove unused import Mock * add test for utils * add dnspython dependency * change fixture name * change fixture name in arguments
1 parent d017528 commit bdd32b4

File tree

6 files changed

+138
-12
lines changed

6 files changed

+138
-12
lines changed

examples/dump.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,10 @@
1010
from loguru import logger
1111

1212
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
13-
from pythclient.pythclient import PythClient, V2_PROGRAM_KEY, V2_FIRST_MAPPING_ACCOUNT_KEY # noqa
13+
from pythclient.pythclient import PythClient # noqa
1414
from pythclient.ratelimit import RateLimit # noqa
1515
from pythclient.pythaccounts import PythPriceAccount # noqa
16+
from pythclient.utils import get_key # noqa
1617

1718
logger.enable("pythclient")
1819

@@ -32,9 +33,11 @@ def set_to_exit(sig: Any, frame: Any):
3233
async def main():
3334
global to_exit
3435
use_program = len(sys.argv) >= 2 and sys.argv[1] == "program"
36+
v2_first_mapping_account_key = get_key("devnet", "mapping")
37+
v2_program_key = get_key("devnet", "program")
3538
async with PythClient(
36-
first_mapping_account_key=V2_FIRST_MAPPING_ACCOUNT_KEY,
37-
program_key=V2_PROGRAM_KEY if use_program else None,
39+
first_mapping_account_key=v2_first_mapping_account_key,
40+
program_key=v2_program_key if use_program else None,
3841
) as c:
3942
await c.refresh_all_prices()
4043
products = await c.get_products()
@@ -57,7 +60,7 @@ async def main():
5760
await ws.connect()
5861
if use_program:
5962
print("Subscribing to program account")
60-
await ws.program_subscribe(V2_PROGRAM_KEY, await c.get_all_accounts())
63+
await ws.program_subscribe(v2_program_key, await c.get_all_accounts())
6164
else:
6265
print("Subscribing to all prices")
6366
for account in all_prices:
@@ -88,7 +91,7 @@ async def main():
8891

8992
print("Unsubscribing...")
9093
if use_program:
91-
await ws.program_unsubscribe(V2_PROGRAM_KEY)
94+
await ws.program_unsubscribe(v2_program_key)
9295
else:
9396
for account in all_prices:
9497
await ws.unsubscribe(account)

pythclient/pythclient.py

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,17 +14,13 @@
1414
from .pythaccounts import PythAccount, PythMappingAccount, PythProductAccount, PythPriceAccount
1515
from . import exceptions, config, ratelimit
1616

17-
V1_FIRST_MAPPING_ACCOUNT_KEY = "ArppEFcsybCLE8CRtQJLQ9tLv2peGmQoKWFuiUWm4KBP"
18-
V2_FIRST_MAPPING_ACCOUNT_KEY = "BmA9Z6FjioHJPpjT39QazZyhDRUdZy2ezwx4GiDdE2u2"
19-
V2_PROGRAM_KEY = "gSbePebfvPy7tRqimPoVecS2UsBvYv46ynrzWocc92s"
20-
2117

2218
class PythClient:
2319
def __init__(self, *,
2420
solana_client: Optional[SolanaClient] = None,
2521
solana_endpoint: str = SOLANA_DEVNET_HTTP_ENDPOINT,
2622
solana_ws_endpoint: str = SOLANA_DEVNET_WS_ENDPOINT,
27-
first_mapping_account_key: str = V1_FIRST_MAPPING_ACCOUNT_KEY,
23+
first_mapping_account_key: str,
2824
program_key: Optional[str] = None,
2925
aiohttp_client_session: Optional[aiohttp.ClientSession] = None) -> None:
3026
self._first_mapping_account_key = SolanaPublicKey(first_mapping_account_key)

pythclient/utils.py

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
import ast
2+
import dns.resolver
3+
from loguru import logger
4+
from typing import Optional
5+
6+
DEFAULT_VERSION = "v2"
7+
8+
9+
# Retrieving keys via DNS TXT records should not be considered secure and is provided as a convenience only.
10+
# Accounts should be stored locally and verified before being used for production.
11+
def get_key(network: str, type: str, version: str = DEFAULT_VERSION) -> Optional[str]:
12+
"""
13+
Get the program or mapping keys from dns TXT records.
14+
15+
Example dns records:
16+
17+
devnet-program-v2.pyth.network
18+
mainnet-program-v2.pyth.network
19+
testnet-mapping-v2.pyth.network
20+
"""
21+
url = f"{network}-{type}-{version}.pyth.network"
22+
try:
23+
answer = dns.resolver.resolve(url, "TXT")
24+
except dns.resolver.NXDOMAIN:
25+
logger.error("TXT record for {} not found", url)
26+
return ""
27+
if len(answer) != 1:
28+
logger.error("Invalid number of records returned for {}!", url)
29+
return ""
30+
# Example of the raw_key:
31+
# "program=FsJ3A3u2vn5cTVofAjvy6y5kwABJAqYWpe4975bi2epH"
32+
raw_key = ast.literal_eval(list(answer)[0].to_text())
33+
# program=FsJ3A3u2vn5cTVofAjvy6y5kwABJAqYWpe4975bi2epH"
34+
_, key = raw_key.split("=", 1)
35+
return key

setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from setuptools import setup
22

3-
requirements = ['aiodns', 'aiohttp>=3.7.4', 'backoff', 'base58', 'flake8', 'loguru']
3+
requirements = ['aiodns', 'aiohttp>=3.7.4', 'backoff', 'base58', 'dnspython', 'flake8', 'loguru']
44

55
setup(
66
name='pythclient',

tests/test_pyth_client.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
_ACCOUNT_HEADER_BYTES, _VERSION_2, PythMappingAccount, PythPriceType, PythProductAccount, PythPriceAccount
77
)
88

9-
from pythclient.pythclient import PythClient, V2_FIRST_MAPPING_ACCOUNT_KEY, V2_PROGRAM_KEY, WatchSession
9+
from pythclient.pythclient import PythClient, WatchSession
1010
from pythclient.solana import (
1111
SolanaClient,
1212
SolanaCommitment,
@@ -24,6 +24,9 @@
2424
# 2) these values are used in get_account_info_resp() and get_program_accounts_resp()
2525
# and so if they are passed in as fixtures, the functions will complain for the args
2626
# while mocking the respective functions
27+
V2_FIRST_MAPPING_ACCOUNT_KEY = 'BmA9Z6FjioHJPpjT39QazZyhDRUdZy2ezwx4GiDdE2u2'
28+
V2_PROGRAM_KEY = 'gSbePebfvPy7tRqimPoVecS2UsBvYv46ynrzWocc92s'
29+
2730
BCH_PRODUCT_ACCOUNT_KEY = '89GseEmvNkzAMMEXcW9oTYzqRPXTsJ3BmNerXmgA1osV'
2831
BCH_PRICE_ACCOUNT_KEY = '4EQrNZYk5KR1RnjyzbaaRbHsv8VqZWzSUtvx58wLsZbj'
2932

tests/test_utils.py

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
from _pytest.logging import LogCaptureFixture
2+
import pytest
3+
4+
from pytest_mock import MockerFixture
5+
6+
from mock import Mock
7+
8+
import dns.resolver
9+
import dns.rdatatype
10+
import dns.rdataclass
11+
import dns.message
12+
import dns.rrset
13+
import dns.flags
14+
15+
from pythclient.utils import get_key
16+
17+
18+
@pytest.fixture()
19+
def answer_program() -> dns.resolver.Answer:
20+
qname = dns.name.Name(labels=(b'devnet-program-v2', b'pyth', b'network', b''))
21+
rdtype = dns.rdatatype.TXT
22+
rdclass = dns.rdataclass.IN
23+
response = dns.message.QueryMessage(id=0)
24+
response.flags = dns.flags.QR
25+
rrset_qn = dns.rrset.from_text(qname, 100, rdclass, rdtype)
26+
rrset_ans = dns.rrset.from_text(qname, 100, rdclass, rdtype, '"program=gSbePebfvPy7tRqimPoVecS2UsBvYv46ynrzWocc92s"')
27+
response.question = [rrset_qn]
28+
response.answer = [rrset_ans]
29+
answer = dns.resolver.Answer(
30+
qname=qname, rdtype=rdtype, rdclass=rdclass, response=response)
31+
answer.rrset = rrset_ans
32+
return answer
33+
34+
35+
@pytest.fixture()
36+
def answer_mapping() -> dns.resolver.Answer:
37+
qname = dns.name.Name(labels=(b'devnet-mapping-v2', b'pyth', b'network', b''))
38+
rdtype = dns.rdatatype.TXT
39+
rdclass = dns.rdataclass.IN
40+
response = dns.message.QueryMessage(id=0)
41+
response.flags = dns.flags.QR
42+
rrset_qn = dns.rrset.from_text(qname, 100, rdclass, rdtype)
43+
rrset_ans = dns.rrset.from_text(qname, 100, rdclass, rdtype, '"mapping=BmA9Z6FjioHJPpjT39QazZyhDRUdZy2ezwx4GiDdE2u2"')
44+
response.question = [rrset_qn]
45+
response.answer = [rrset_ans]
46+
answer = dns.resolver.Answer(
47+
qname=qname, rdtype=rdtype, rdclass=rdclass, response=response)
48+
answer.rrset = rrset_ans
49+
return answer
50+
51+
52+
@pytest.fixture()
53+
def mock_dns_resolver_resolve(mocker: MockerFixture) -> Mock:
54+
mock = Mock()
55+
mocker.patch('dns.resolver.resolve', side_effect=mock)
56+
return mock
57+
58+
59+
def test_utils_get_program_key(mock_dns_resolver_resolve: Mock, answer_program: dns.resolver.Answer) -> None:
60+
mock_dns_resolver_resolve.return_value = answer_program
61+
program_key = get_key("devnet", "program")
62+
assert program_key == "gSbePebfvPy7tRqimPoVecS2UsBvYv46ynrzWocc92s"
63+
64+
65+
def test_utils_get_mapping_key(mock_dns_resolver_resolve: Mock, answer_mapping: dns.resolver.Answer) -> None:
66+
mock_dns_resolver_resolve.return_value = answer_mapping
67+
mapping_key = get_key("devnet", "mapping")
68+
assert mapping_key == "BmA9Z6FjioHJPpjT39QazZyhDRUdZy2ezwx4GiDdE2u2"
69+
70+
71+
def test_utils_get_mapping_key_not_found(mock_dns_resolver_resolve: Mock,
72+
answer_mapping: dns.resolver.Answer,
73+
caplog: LogCaptureFixture) -> None:
74+
mock_dns_resolver_resolve.side_effect = dns.resolver.NXDOMAIN
75+
exc_message = f'TXT record for {str(answer_mapping.response.canonical_name())[:-1]} not found'
76+
key = get_key("devnet", "mapping")
77+
assert exc_message in caplog.text
78+
assert key == ""
79+
80+
81+
def test_utils_get_mapping_key_invalid_number(mock_dns_resolver_resolve: Mock,
82+
answer_mapping: dns.resolver.Answer,
83+
caplog: LogCaptureFixture) -> None:
84+
answer_mapping.rrset = None
85+
mock_dns_resolver_resolve.return_value = answer_mapping
86+
exc_message = f'Invalid number of records returned for {str(answer_mapping.response.canonical_name())[:-1]}!'
87+
key = get_key("devnet", "mapping")
88+
assert exc_message in caplog.text
89+
assert key == ""

0 commit comments

Comments
 (0)