diff --git a/src/adguardhome/adguardhome.py b/src/adguardhome/adguardhome.py index 345f02eb..38859d7f 100644 --- a/src/adguardhome/adguardhome.py +++ b/src/adguardhome/adguardhome.py @@ -14,6 +14,7 @@ from .filtering import AdGuardHomeFiltering from .parental import AdGuardHomeParental from .querylog import AdGuardHomeQueryLog +from .rewrite import AdGuardHomeRewrite from .safebrowsing import AdGuardHomeSafeBrowsing from .safesearch import AdGuardHomeSafeSearch from .stats import AdGuardHomeStats @@ -77,6 +78,7 @@ def __init__( # noqa: PLR0913 self.filtering = AdGuardHomeFiltering(self) self.parental = AdGuardHomeParental(self) self.querylog = AdGuardHomeQueryLog(self) + self.rewrite = AdGuardHomeRewrite(self) self.safebrowsing = AdGuardHomeSafeBrowsing(self) self.safesearch = AdGuardHomeSafeSearch(self) self.stats = AdGuardHomeStats(self) @@ -90,7 +92,7 @@ async def request( data: Any | None = None, json_data: dict[str, Any] | None = None, params: Mapping[str, str] | None = None, - ) -> dict[str, Any]: + ) -> Any: """Handle a request to the AdGuard Home instance. Make a request against the AdGuard Home API and handles the response. @@ -107,7 +109,7 @@ async def request( ------- The response from the API. In case the response is a JSON response, the method will return a decoded JSON response as a Python - dictionary. In other cases, it will return the RAW text response. + dictionary or list. In other cases, it will return the RAW text response. Raises: ------ diff --git a/src/adguardhome/rewrite.py b/src/adguardhome/rewrite.py new file mode 100644 index 00000000..9f6ce9da --- /dev/null +++ b/src/adguardhome/rewrite.py @@ -0,0 +1,74 @@ +"""Asynchronous Python client for the AdGuard Home API.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from .exceptions import AdGuardHomeError + +if TYPE_CHECKING: + from . import AdGuardHome + + +@dataclass +class AdGuardHomeRewrite: + """Controls AdGuard Home DNS rewrites.""" + + adguard: AdGuardHome + + async def list(self) -> list[dict[str, str]]: + """Return all defined DNS rewrites. + + Returns: + A list of all defined DNS rewrites in the AdGuardHome instance. + Each rewrite is represented as a dictionary with 'domain' and 'answer' keys. + + """ + return await self.adguard.request("rewrite/list") + + async def add(self, domain: str, answer: str) -> None: + """Add a new DNS rewrite rule to AdGuard Home. + + Args: + ---- + domain: The domain pattern to rewrite (e.g., "*.example.com"). + answer: The IP address or domain to rewrite to. + + Raises: + ------ + AdGuardHomeError: Failed adding the DNS rewrite rule. + + """ + try: + await self.adguard.request( + "rewrite/add", + method="POST", + json_data={"domain": domain, "answer": answer}, + ) + except AdGuardHomeError as exception: + msg = "Failed to add DNS rewrite rule to AdGuard Home" + raise AdGuardHomeError(msg) from exception + + async def delete(self, domain: str, answer: str) -> None: + """Delete a DNS rewrite rule from AdGuard Home. + + Args: + ---- + domain: The domain pattern of the rewrite rule to delete. + answer: The IP address or domain of the rewrite rule to delete. + + Raises: + ------ + AdGuardHomeError: Failed to delete DNS rewrite rule. + + """ + try: + await self.adguard.request( + "rewrite/delete", + method="POST", + json_data={"domain": domain, "answer": answer}, + ) + except AdGuardHomeError as exception: + msg = "Failed to delete DNS rewrite rule from AdGuard Home" + raise AdGuardHomeError(msg) from exception diff --git a/tests/test_rewrite.py b/tests/test_rewrite.py new file mode 100644 index 00000000..3a70119a --- /dev/null +++ b/tests/test_rewrite.py @@ -0,0 +1,98 @@ +"""Tests for `adguardhome.rewrite`.""" + +import aiohttp +import pytest +from aresponses import Response, ResponsesMockServer + +from adguardhome import AdGuardHome +from adguardhome.exceptions import AdGuardHomeError + + +async def test_list(aresponses: ResponsesMockServer) -> None: + """Test getting all DNS rewrite rules from AdGuard Home rewrite.""" + aresponses.add( + "example.com:3000", + "/control/rewrite/list", + "GET", + aresponses.Response( + status=200, + headers={"Content-Type": "application/json"}, + text='[{"domain": "*.example.com", "answer": "192.168.1.2"},' + '{"domain": "*.example.com", "answer": "192.168.1.2"}]', + ), + ) + + aresponses.add( + "example.com:3000", + "/control/rewrite/list", + "GET", + aresponses.Response( + status=200, headers={"Content-Type": "application/json"}, text="[]" + ), + ) + + async with aiohttp.ClientSession() as session: + adguard = AdGuardHome("example.com", session=session) + result = await adguard.rewrite.list() + assert result == [ + {"domain": "*.example.com", "answer": "192.168.1.2"}, + {"domain": "*.example.com", "answer": "192.168.1.2"}, + ] + result = await adguard.rewrite.list() + assert result == [] + + +async def test_add(aresponses: ResponsesMockServer) -> None: + """Test add new DNS rewrite to AdGuard rewrite.""" + + async def response_handler(request: aiohttp.ClientResponse) -> Response: + data = await request.json() + assert data == {"domain": "*.example.com", "answer": "192.168.1.2"} + return aresponses.Response(status=200) + + aresponses.add( + "example.com:3000", + "/control/rewrite/add", + "POST", + response_handler, + ) + aresponses.add( + "example.com:3000", + "/control/rewrite/add", + "POST", + aresponses.Response(status=400, text="Bad Request"), + ) + + async with aiohttp.ClientSession() as session: + adguard = AdGuardHome("example.com", session=session) + await adguard.rewrite.add("*.example.com", "192.168.1.2") + with pytest.raises(AdGuardHomeError): + await adguard.rewrite.add("*.example.com", "192.168.1.2") + + +async def test_remove(aresponses: ResponsesMockServer) -> None: + """Test deleting DNS rewrite from AdGuard Home rewrite.""" + + async def response_handler(request: aiohttp.ClientResponse) -> Response: + data = await request.json() + assert data == {"domain": "*.example.com", "answer": "192.168.1.2"} + return aresponses.Response(status=200) + + aresponses.add( + "example.com:3000", + "/control/rewrite/delete", + "POST", + response_handler, + ) + aresponses.add( + "example.com:3000", + "/control/rewrite/delete", + "POST", + aresponses.Response(status=400, text="Bad Request"), + ) + + async with aiohttp.ClientSession() as session: + adguard = AdGuardHome("example.com", session=session) + await adguard.rewrite.delete("*.example.com", "192.168.1.2") + with pytest.raises(AdGuardHomeError): + await adguard.rewrite.delete("*.example.com", "192.168.1.2")