134 lines
6.2 KiB
Python
134 lines
6.2 KiB
Python
import re
|
|
|
|
# Read the server.py file
|
|
with open('src/vultr_mcp/server.py', 'r') as f:
|
|
content = f.read()
|
|
|
|
# Add DNS tool handling to call_tool function
|
|
# Find the list_os handling and insert DNS handling after it
|
|
list_os_pattern = r"(elif request\.name == \"list_os\":\s*result = vultr_client\.list_os\(\)\s*)"
|
|
match = re.search(list_os_pattern, content)
|
|
if match:
|
|
list_os_section = match.group(1)
|
|
dns_handling = ''' elif request.name == "list_dns_domains":
|
|
result = vultr_client.list_dns_domains()
|
|
|
|
elif request.name == "create_dns_domain":
|
|
domain = request.arguments.get("domain")
|
|
ip = request.arguments.get("ip")
|
|
if not domain or not ip:
|
|
raise ValueError("domain and ip are required")
|
|
result = vultr_client.create_dns_domain(domain, ip)
|
|
|
|
elif request.name == "delete_dns_domain":
|
|
domain = request.arguments.get("domain")
|
|
if not domain:
|
|
raise ValueError("domain is required")
|
|
result = vultr_client.delete_dns_domain(domain)
|
|
|
|
elif request.name == "get_dns_domain":
|
|
domain = request.arguments.get("domain")
|
|
if not domain:
|
|
raise ValueError("domain is required")
|
|
result = vultr_client.get_dns_domain(domain)
|
|
|
|
elif request.name == "list_dns_records":
|
|
domain = request.arguments.get("domain")
|
|
if not domain:
|
|
raise ValueError("domain is required")
|
|
result = vultr_client.list_dns_records(domain)
|
|
|
|
elif request.name == "create_dns_record":
|
|
domain = request.arguments.get("domain")
|
|
record_type = request.arguments.get("type")
|
|
name = request.arguments.get("name")
|
|
data = request.arguments.get("data")
|
|
ttl = request.arguments.get("ttl", 300)
|
|
priority = request.arguments.get("priority", 0)
|
|
if not domain or not record_type or not name or not data:
|
|
raise ValueError("domain, type, name, and data are required")
|
|
result = vultr_client.create_dns_record(domain, record_type, name, data, ttl, priority)
|
|
|
|
elif request.name == "update_dns_record":
|
|
domain = request.arguments.get("domain")
|
|
record_id = request.arguments.get("record_id")
|
|
if not domain or not record_id:
|
|
raise ValueError("domain and record_id are required")
|
|
# Filter out None values
|
|
update_args = {k: v for k, v in request.arguments.items()
|
|
if k not in ["domain", "record_id"] and v is not None}
|
|
result = vultr_client.update_dns_record(domain, record_id, **update_args)
|
|
|
|
elif request.name == "delete_dns_record":
|
|
domain = request.arguments.get("domain")
|
|
record_id = request.arguments.get("record_id")
|
|
if not domain or not record_id:
|
|
raise ValueError("domain and record_id are required")
|
|
result = vultr_client.delete_dns_record(domain, record_id)
|
|
|
|
'''
|
|
# Insert DNS handling after list_os
|
|
new_content = content.replace(list_os_section, list_os_section + dns_handling)
|
|
content = new_content
|
|
|
|
# Add DNS resources to list_resources function
|
|
resources_pattern = r"(Resource\(\s*uri=\"vultr://os/list\",\s*name=\"Vultr Operating Systems\",\s*description=\"List of available Vultr operating systems\"\s*\)\s*,?\s*)"
|
|
match = re.search(resources_pattern, content, re.DOTALL)
|
|
if match:
|
|
os_resource = match.group(1)
|
|
dns_resources = ''' Resource(
|
|
uri="vultr://dns/domains",
|
|
name="Vultr DNS Domains",
|
|
description="List of DNS domains"
|
|
),
|
|
Resource(
|
|
uri="vultr://dns/records",
|
|
name="Vultr DNS Records",
|
|
description="DNS records for a domain"
|
|
),
|
|
'''
|
|
new_content = content.replace(os_resource, os_resource + dns_resources)
|
|
content = new_content
|
|
|
|
# Add DNS prompts to list_prompts function
|
|
prompts_pattern = r"(Prompt\(\s*name=\"monitor_resources\",\s*description=\"Monitor Vultr resources\",\s*arguments=\[\s*\{.*?\}\s*\]\s*\)\s*,?\s*)"
|
|
match = re.search(prompts_pattern, content, re.DOTALL)
|
|
if match:
|
|
monitor_prompt = match.group(1)
|
|
dns_prompts = ''' Prompt(
|
|
name="manage_dns",
|
|
description="Manage DNS domains and records",
|
|
arguments=[
|
|
{"name": "action", "description": "Action to perform (create_domain, delete_domain, list_records, create_record, update_record, delete_record)", "required": True},
|
|
{"name": "domain", "description": "Domain name", "required": True},
|
|
{"name": "record_id", "description": "Record ID (for update/delete)", "required": False},
|
|
{"name": "record_type", "description": "Record type (A, AAAA, CNAME, etc.)", "required": False},
|
|
{"name": "name", "description": "Record name", "required": False},
|
|
{"name": "data", "description": "Record data", "required": False},
|
|
{"name": "ip", "description": "IP address (for create domain)", "required": False},
|
|
]
|
|
),
|
|
'''
|
|
new_content = content.replace(monitor_prompt, monitor_prompt + dns_prompts)
|
|
content = new_content
|
|
|
|
# Add DNS resource handling to get_resource function
|
|
get_resource_pattern = r"(elif resource_uri == \"vultr://os/list\":\s*result = vultr_client\.list_os\(\)\s*)"
|
|
match = re.search(get_resource_pattern, content)
|
|
if match:
|
|
os_resource_handling = match.group(1)
|
|
dns_resource_handling = ''' elif resource_uri == "vultr://dns/domains":
|
|
result = vultr_client.list_dns_domains()
|
|
elif resource_uri == "vultr://dns/records":
|
|
# This requires a domain parameter, so we'll return instructions
|
|
result = {"message": "Please use the list_dns_records tool with domain parameter to get DNS records"}
|
|
'''
|
|
new_content = content.replace(os_resource_handling, os_resource_handling + dns_resource_handling)
|
|
content = new_content
|
|
|
|
# Write the updated content back
|
|
with open('src/vultr_mcp/server.py', 'w') as f:
|
|
f.write(content)
|
|
|
|
print("Server.py updated successfully with DNS functionality")
|