<commit_before>from twisted.internet import reactor
from twisted.names import dns, client, server
from rpyc.utils.factory import ssh_connect
from plumbum import SshMachine

class DNSServerFactory(server.DNSServerFactory):
    def handleQuery(self, message, protocol, address):
        try:
            query = message.queries[0]
            target = query.name.name
            print 'Target:', target

            query_type = target.split('.')[0]
            if query_type == 'ns1':
                A = dns.RRHeader(name=target, type=dns.A, cls=dns.IN, ttl=0,
                                payload=dns.Record_A(address='192.168.3.11', ttl=None))
                args = (self, ([A], [], []), protocol, message, address)
                return server.DNSServerFactory.gotResolverResponse(*args)
            elif query_type == 'ns2':
                A = dns.RRHeader(name=target, type=dns.A, cls=dns.IN, ttl=0,
                                payload=dns.Record_A(address='172.16.17.32', ttl=None))
                args = (self, ([A], [], []), protocol, message, address)
                return server.DNSServerFactory.gotResolverResponse(*args)
            elif query_type == 'full':
                origin = target.split('.')[3].split('---')
                origin_ns_name = '.'.join(origin[4:])
                origin_ip = '.'.join(origin[:4])
                target = '.'.join(target.split('.')[3:])
                node_ip = '.'.join(target.split('.')[1].split('---'))
                print 'Full', query_type, origin_ip, origin_ns_name, node_ip
            else:
                origin = target.split('.')[2].split('---')
                origin_ns_name = '.'.join(origin[4:])
                origin_ip = '.'.join(origin[:4])
                target = '.'.join(target.split('.')[2:])
                print query_type, origin_ip, origin_ns_name

                NS = dns.RRHeader(name=target, type=dns.NS, cls=dns.IN, ttl=0, auth=True,
                                 payload=dns.Record_NS(name=origin_ns_name, ttl=0))
                A = dns.RRHeader(name=origin_ns_name, type=dns.A, cls=dns.IN, ttl=0,
                                payload=dns.Record_A(address=origin_ip, ttl=None))

                ans = []
                auth = [NS]
                add = [A]
                args = (self, (ans, auth, add), protocol, message, address)

                return server.DNSServerFactory.gotResolverResponse(*args)
        except Exception, e:
            print "Bad Request", e

factory = DNSServerFactory()
protocol = dns.DNSDatagramProtocol(factory)

reactor.listenUDP(53, protocol)
reactor.listenTCP(53, factory)
reactor.run()<commit_msg>added code to make RPC from central name server<commit_after>from twisted.internet import reactor
from twisted.names import dns, client, server
from rpyc.utils.factory import ssh_connect
from plumbum import SshMachine
import argparse

parser = argparse.ArgumentParser(description='Central Name Server')
parser.add_argument('--full', default=False, action='store_true', help='This instance will act as the endpoint for integration testing')

arguments = parser.parse_args()

class DNSServerFactory(server.DNSServerFactory):
    def handleQuery(self, message, protocol, address):
        try:
            query = message.queries[0]
            target = query.name.name
            print 'Target:', target

            query_type = target.split('.')[0]
            if query_type == 'ns1':
                A = dns.RRHeader(name=target, type=dns.A, cls=dns.IN, ttl=0,
                                payload=dns.Record_A(address='192.168.3.11', ttl=None))
                args = (self, ([A], [], []), protocol, message, address)
                return server.DNSServerFactory.gotResolverResponse(*args)
            elif query_type == 'ns2':
                A = dns.RRHeader(name=target, type=dns.A, cls=dns.IN, ttl=0,
                                payload=dns.Record_A(address='172.16.17.32', ttl=None))
                args = (self, ([A], [], []), protocol, message, address)
                return server.DNSServerFactory.gotResolverResponse(*args)
            elif query_type == 'full':
                query_id = int(target.split('.')[2])
                origin = target.split('.')[3].split('---')
                origin_ns_name = '.'.join(origin[4:])
                origin_ip = '.'.join(origin[:4])
                target = '.'.join(target.split('.')[3:])
                node_ip = '.'.join(target.split('.')[1].split('---'))
                print 'Full', query_type, origin_ip, origin_ns_name, node_ip
                if arguments.full:
                    rem = SshMachine(host, user='ucb_268_measure', keyfile='~/.ssh/id_rsa')
                    conn = ssh_connect(rem, 18861)
                    conn.root.exposed_full_response(query_id, 'End Point Reached')
            else:
                origin = target.split('.')[2].split('---')
                origin_ns_name = '.'.join(origin[4:])
                origin_ip = '.'.join(origin[:4])
                target = '.'.join(target.split('.')[2:])
                print query_type, origin_ip, origin_ns_name

                NS = dns.RRHeader(name=target, type=dns.NS, cls=dns.IN, ttl=0, auth=True,
                                 payload=dns.Record_NS(name=origin_ns_name, ttl=0))
                A = dns.RRHeader(name=origin_ns_name, type=dns.A, cls=dns.IN, ttl=0,
                                payload=dns.Record_A(address=origin_ip, ttl=None))

                ans = []
                auth = [NS]
                add = [A]
                args = (self, (ans, auth, add), protocol, message, address)

                return server.DNSServerFactory.gotResolverResponse(*args)
        except Exception, e:
            print "Bad Request", e

factory = DNSServerFactory()
protocol = dns.DNSDatagramProtocol(factory)

reactor.listenUDP(53, protocol)
reactor.listenTCP(53, factory)
reactor.run()
{'fraction_non_alphanumeric': 0.1404370658986956, 'fraction_numerical': 0.015077079451126546, 'mean_word_length': 1.9984763839512443, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 4, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3699086', 'n_tokens_mistral': 1869, 'n_tokens_neox': 1751, 'n_words': 341}