import socket
from shutil import copyfile, move
from os import remove
from os.path import isfile


class MookDns:
    # Monkey patch to force IPv4 resolution
    def __init__(self, ip):
        self.ip = ip

    def __enter__(self):
        self.old_getaddrinfo = socket.getaddrinfo
        def new_getaddrinfo(*args, **kwargs):
            ret = self.old_getaddrinfo(*args, **kwargs)
            dns = list(ret[0])
            dns[-1] = (self.ip, dns[-1][1])
            return [dns]
        socket.getaddrinfo = new_getaddrinfo
        return self

    def __exit__(self, exc_type, exc, tb):
        socket.getaddrinfo = self.old_getaddrinfo


class MookDnsSystem:
    # Monkey patch to force IPv4 resolution
    def __init__(self, dns, ip):
        self.dns = dns
        self.ip = ip

    def __enter__(self):
        if not isfile('/etc/hosts.risotto'):
            copyfile('/etc/hosts', '/etc/hosts.risotto')
        with open('/etc/hosts.risotto', 'r') as risotto:
            with open('/etc/hosts', 'w') as hosts:
                for line in risotto.readlines():
                    if self.dns not in line:
                        hosts.write(line)
                hosts.write(f'{self.ip}  {self.dns}')

    def __exit__(self, exc_type, exc, tb):
        remove('/etc/hosts')
        move('/etc/hosts.risotto', '/etc/hosts')