forked from stove/dataset
45 lines
1.3 KiB
Python
45 lines
1.3 KiB
Python
|
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')
|