dnscode/main.py

from enum import Enum
from dataclasses import dataclass, field
import time
import ipaddress
import dns
import fqdn

class InvaliadDataException(Exception):
	"""Exception raied when invaliad data is passed to a record"""

	def __init__(self, message):
		self.message	= message
		super().__init__(self, message)

@dataclass
class Record:
	def __int__(self, name: str = '@', ttl: str = 3600, rtype: str = 'A', data: str = '0.0.0.0'):
		self.rtype: str	= rtype
		self.name	= name
		self.data	= data
		self.ttl	= ttl

	def __str__(self):
		return f"{self.name} {self.ttl} {self.rclass} {self.rtype} {self.data}"

	rclass: str = 'IN'
	rtype: str	= 'A'
	name: str	= '@'
	data: str	= '0.0.0.0'
	ttl: int	= 6400

@dataclass
class A(Record):
	def __init__(self, name: str = '@', ttl: str = 3600, target: str = '0.0.0.0'):
		if isinstance(ipaddress.ip_address(data), ipaddress.IPv4Address):
			self.data = data
		else:
			raise InvaliadDataException(message=f'{data} is not a valid IPv4 Address.')

		self.rtype	= 'A'
		self.name	= name
		self.ttl	= ttl

@dataclass
class AAAA(Record):
	def __init__(self, name: str = '@', ttl: str = 3600, target: str = '0.0.0.0'):
		if isinstance(ipaddress.ip_address(data), ipaddress.IPv4Address):
			self.data = data
		else:
			raise InvaliadDataException(message=f'{data} is not a valiad IPv6 Address.')

		self.rtype	= 'AAAA'
		self.name	= name
		self.ttl	= ttl

@dataclass
class CNAME(Record):
	def __init__(self, name: str = '@', ttl: str = 3600, target: str = '0.0.0.0'):
		self.rtype	= 'CNAME'
		self.name	= name
		self.ttl	= ttl

		if(fqdn.FQDN(target).is_valid):
			self.data = target
		else:
			raise InvaliadDataException

@dataclass
class MX(Record):
	def __init__(self, name: str = '@', ttl: str = 3600, priority: int = '10', host: str = 'example.com'):
		self.rtype	= 'MX'
		self.name	= name
		self.ttl	= ttl
		self.priority	= priority
		self.data	= "{ttl} {priority}"

		if(fqdn.FQDN(host).is_valid):
			self.host = host
		else:
			raise InvaliadDataException(message='{host} is not a valid FQDN')

@dataclass
class NS(Record):
	def __init__(self, name: str = '@', ttl: str = 3600, target: str = 'example.com'):
		self.rtype	= 'MX'
		self.name	= name
		self.ttl	= ttl

		if (fqdn.FQDN(host).is_valid):
			self.host = host
		else:
			raise InvaliadDataException(message='{target} is nod a valid FQDN')
		self.data	= target

@dataclass
class PTR(Record):
	def __init__(self, name: str = '@', ttl: str = 3600, host: str = '0.0.0.0'):
		self.rtype	= 'PTR'
		self.name	= name
		self.ttl	= ttl

		if(fqdn.FQDN(host).is_valid):
			self.data = host
		else:
			raise InvaliadDataException(message='{host} is not a valid FQDN')

# TODO: Cleanup. I have no idea why I have _str_() defined.
@dataclass
class SOA(Record):
	def __init__(self, name: str = '@', mname: str = 'ns1.example.com', rname: str = 'admin.example.com',
				serial: int = int(time.time()), refresh: int = 86400, retry: int = 7200,
				expire: int = 15552000, ttl: int = 21700
		):
		self.rtype		= "SOA"
		self.name		= name
		self.mname		= mname
		self.rname		= rname
		self.serial		= serial
		self.refresh	= refresh
		self.retry		= retry
		self.expire		= expire
		self.ttl		= ttl
		self.data		= "{name} {ttl} IN SOA {mname} {rname} {serial} {refresh} {retry} {expire} {ttl}"

@dataclass
class SRV(Record):
	def __int__(self, name: str = '@', ttl: str = 3600, service: str = "service", protocol: str = 'proto',
				priority: int = 10, weight: int = 10, port: int = 0, target: str = target
				):
		self.rtype	= 'PTR'
		self.name	= '_{service}._{protocol}.name'
		self.ttl 	= ttl
		self.service	= service
		self.protocol	= protocol
		self.priority	= priority
		self.weight		= weight
		self.port	= port
		if (fqdn.FQDN(host).is_valid):
			self.target	= target
		else:
			raise InvaliadDataException(message='{target} is not a valiad FQDN')
		self.data	= "{priority} {weight} {port} {target}"

@dataclass
class Zone:
	origin: str
	records: list = field(default_factory=list)

	def __str__(self):
		zone: str = ''

		for record in self.records:
			zone.join(record)

		return zone

	def __mkfqdn(self, name):
		if name[-1] != '.':
			return name + '.' + self.origin
		else:
			return name

	def new_A(self, name: str = '@', ttl: int = 3600, data: str = '0.0.0.0'):
		name = self.__mkfqdn(name)
		self.add(A(name=name, ttl=ttl, data=data))

	def new_AAAA(self, name: str = '@', ttl: int = 3600, data: str = '0.0.0.0'):
		name = self.__mkfqdn(name)
		self.add(AAAA(name=name, ttl=ttl, data=data))

	def new_CNAME(self, name: str = '@', ttl: int = 3600, data: str = '0.0.0.0'):
		name = self.__mkfqdn(name)
		self.add(CNAME(name=name, ttl=ttl, data=data))

	def new_MX(self, name: str = '@', ttl: int = 3600, priority: int = 10, host: str = 'example.com'):
		name = self.__mkfqdn(name)
		self.add(MX(name=name, ttl=ttl, priority=priority, host=host))

	def new_soa(self, mname: str = 'ns1.example.com', rname: str = 'admin.example.com', serial: int = int(time.time()), refresh: int = 86400, retry: int = 7200, expire: int = 15552000, ttl: int = 21700):
		mname = self.__mkfqdn(name)
		self.add(SOA(mname=mname, rname=rname, serial=serial, refresh=refresh, retry=retry, expire=expire, ttl=ttl))

	def new_record(self, name: str = '@', ttl: int = 3600, rtype: str = 'A', data: str = '0.0.0.0'):
		name = self.__mkfqdn(name)
		self.add(name=name, ttl=ttl, rtype=rtype, data=data)

	def add(self, record: Record):
		self.records.append(record)

	def save_file(self, filepath: str):
		with open(filepath, 'w') as file:
			for record in self.records:
				file.write(str(record) + '\n')
				print(str(record))
		file.close()