from langchain_core.pydantic_v1 import BaseModel, Field, root_validator
from langchain.tools import tool
import dns.resolver, dns.reversename
import validators

"""
Tools to be run by my custom agent.

So far, these are the same tooling provided by the example in https://github.com/wu4f/cs410g-src/blob/main/04_Agents/07_tools_custom_agent.py,
But I will be updating this to build my own tooling
"""

class LookupNameInput(BaseModel):
    hostname: str = Field(description="Should be a hostname such as www.google.com")
    @root_validator
    def is_dns_address(cls, values: dict[str,any]) -> str:
        if validators.domain(values.get("hostname")):
            return values
        raise ValueError("Malformed hostname")
class LookupIPInput(BaseModel):
    address: str = Field(description="Should be an IP address such as 208.91.197.27 or 143.95.239.83")
    @root_validator
    def is_ip_address(cls, values: dict[str,any]) -> str:
        if validators.ip_address.ipv4(values.get("address")):
            return values
        raise ValueError("Malformed IP address")
    

@tool("lookup_name",args_schema=LookupNameInput, return_direct=False)
def lookup_name(hostname):
    """Given a DNS hostname, it will return its IPv4 addresses"""
    result = dns.resolver.resolve(hostname, 'A')
    res = [ r.to_text() for r in result ]
    return res[0]


@tool("lookup_ip", args_schema=LookupIPInput, return_direct=False)
def lookup_ip(address):
    """Given an IP address, returns names associated with it"""
    n = dns.reversename.from_address(address)
    result = dns.resolver.resolve(n, 'PTR')
    res = [ r.to_text() for r in result ]
    return res[0]