44 lines
1.6 KiB
Python
44 lines
1.6 KiB
Python
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]
|