#!/usr/bin/env python3

import getpass
import requests
import time
from bs4 import BeautifulSoup

default_user = 'SecT0uch'

s = requests.Session()


def login(user):
    print('Logging in...')

    api_login_endpoint = 'https://api.www.root-me.org/login/'

    pw = getpass.getpass()

    # Send the login request
    params = {'login': user, 'password': pw}
    r = s.get(api_login_endpoint, params=params)

    # Get the session cookie
    token = r.json()[0]['info']['spip_session']

    # Set the session cookie as default for next requests
    cookie = {'spip_session': token}
    s.cookies.update(cookie)


def col_print(lines, term_width=None, indent=0, pad=2):
    if not term_width:
        import shutil
        size = shutil.get_terminal_size((80, 20))
        term_width = size.columns
    n_lines = len(lines)
    if n_lines == 0:
        return

    col_width = max(len(line) for line in lines)
    n_cols = int((term_width + pad - indent)/(col_width + pad))
    n_cols = min(n_lines, max(1, n_cols))

    col_len = int(n_lines/n_cols) + (0 if n_lines % n_cols == 0 else 1)
    if (n_cols - 1) * col_len >= n_lines:
        n_cols -= 1

    cols = [lines[i*col_len: i*col_len + col_len] for i in range(n_cols)]

    rows = list(zip(*cols))
    rows_missed = zip(*[col[len(rows):] for col in cols[:-1]])
    rows.extend(rows_missed)

    for row in rows:
        print(" "*indent + (" "*pad).join(line.ljust(col_width) for line in row))


def is_ctf_started(user):
    print('Checking if a CTF is already started...')

    url = 'https://www.root-me.org/fr/Capture-The-Flag/CTF-all-the-day/'
    r = s.get(url)
    soup = BeautifulSoup(r.text, 'html.parser')

    for room in soup.find_all('tr', class_=['row_even', 'row_odd']):
        room_users = room.find_all('a', class_='forum')
        for room_user in room_users:
            room_user = room_user.string
            if room_user == user:
                room_name = room.find('a').string
                room_id = room.find('a').get('href').split('id_salle=')[1].split('&')[0]
                print(f' ==> You are registered in {room_name}')
                get_ctf_info(room_id)
                return room_id, room_name
    print(' ==> You are not registered in any room yet.')


def get_ctf_list():
    print('Getting CTF list...')

    api_ctf_endpoint = 'https://api.www.root-me.org/environnements_virtuels/'
    r = s.get(api_ctf_endpoint)
    ctf_list = []

    while True:  # For each page
        results = r.json()
        ctfs = results[0]
        for key in ctfs:
            ctf = ctfs[key]
            id = ctf['id_environnement_virtuel'].ljust(3, ' ')
            name = ctf['nom']

            ctf_list.append(f"{id} - {name}")

        # Check if next page
        type = results[-1]['rel']
        if type == 'previous':  # If last page break
            col_print(ctf_list)
            break
        elif type == "next":
            next_page = results[-1]['href']
            time.sleep(0.5)   # Avoid getting HTTP 429 (Too Many Requests)
            r = s.get(next_page)


def start_ctf(ctf_id):
    def get_free_room():
        print('Finding a free room...')

        url = 'https://www.root-me.org/fr/Capture-The-Flag/CTF-all-the-day/'
        r = s.get(url)
        soup = BeautifulSoup(r.text, 'html.parser')

        for room in soup.find_all('tr', class_=['row_even', 'row_odd']):
            if room.find('img', alt='waiting'):
                freeroom_name = room.find('a').string
                freeroom_id = room.find('a').get('href').split('id_salle=')[1].split('&')[0]
                print(f"  ==> Room {freeroom_name} is free!")
                return freeroom_name, freeroom_id

    freeroom_name, freeroom_id = get_free_room()

    params = {'lang': 'en', 'page': 'ctf_alltheday', 'id_salle': freeroom_id}
    r = s.get('https://www.root-me.org/', params=params)
    soup = BeautifulSoup(r.text, 'html.parser')

    form = soup.find('form', {'name': 'formulaire_ctf_alltheday_affiche_partie'})
    formulaire_action_args = form.find('input', {'name': 'formulaire_action_args'})['value']
    ctf_name = form.find('option', {'value': ctf_id}).string

    data = {'var_ajax': 'form',
            'page': 'ctf_alltheday',
            'lang': 'en',
            'formulaire_action': 'ctf_alltheday_affiche_partie',
            'formulaire_action_args': formulaire_action_args,
            'id_environnement_virtuel': ctf_id,
            'action_partie': 'enregistrer_choix',
            'enregistrer_choix': 'Save',
            'demarrer': 'Start+the+game'}

    s.post('https://www.root-me.org/', params=params, data=data)
    print(f'Starting "{ctf_name}" on {freeroom_name}...')
    print('Read the description of the CTF meanwhile:')
    get_ctf_info(freeroom_id)

    while 'Waiting for room to start':
        time.sleep(5)
        r = s.get('https://www.root-me.org/', params=params)
        soup = BeautifulSoup(r.text, 'html.parser')
        success = soup.find('div', class_="success")
        if success is not None:
            print(f"The environnement is now available at {freeroom_name}.root-me.org")
            return freeroom_name


def get_ctf_info(room_id):
    params = {'lang': 'en', 'page': 'ctf_alltheday', 'id_salle': room_id}
    r = s.get('https://www.root-me.org/', params=params)
    soup = BeautifulSoup(r.text, 'html.parser')
    desc = soup.find('ul', class_='spip')
    # desc = re.sub(r'(.*)(Description.*)(Game duration.*min)', r'\1\n\2\n\n\3\n', desc, flags=re.DOTALL)
    desc = str(desc).replace('<p>', '\n').replace('</p>', '\n')
    soup = BeautifulSoup(desc, 'html.parser')
    desc = soup.text
    print(f'{"#" * 30}\n{desc}\n{"#" * 30}')


def yes_or_no(question, default=None):
    while "The answer is invalid":
        reply = input(question).lower().strip()
        if reply[:1] == 'y':
            return True
        if reply[:1] == 'n':
            return False
        if not reply:
            if default == 'Y':
                return True
            elif default == 'N':
                return False


def config_ctf(room_name):
    import os
    import subprocess

    user = input('SSH user: ')
    pw_orig = getpass.getpass()
    host = f"{room_name}.root-me.org"

    def send_ssh_key():
        if "ssh-copy-id" not in os.listdir("/usr/bin") or "sshpass" not in os.listdir("/usr/bin"):
            import sys
            print("ssh-copy-id and sshpass are required.\nAfter installation, be sure to have SSH keys.")
            sys.exit()

        command = f"sshpass -p {pw_orig} ssh-copy-id {user}@{host}"
        subprocess.call(command, shell=True)

    def change_password():
        import secrets
        import string

        alphabet = string.ascii_letters + string.digits
        pw = ''.join(secrets.choice(alphabet) for i in range(20))
        local_command = f'echo "{pw_orig}\n{pw}\n{pw}" | passwd'
        command = f"ssh {user}@{host} '{local_command}'"
        subprocess.call(command, shell=True)
        print(f"Password changed: {pw}")

    def send_kitty_terminfo():
        """For kitty shell users"""
        command = f'infocmp xterm-kitty | ssh {user}@{host} tic -x -o \\~/.terminfo /dev/stdin'
        subprocess.call(command, shell=True, stdout=subprocess.DEVNULL)

    send_ssh_key()
    change_password()

    if "kitty" in os.listdir("/usr/bin"):
        send_kitty_terminfo()

    print(f'The environnement is now configured, you can connect with: ssh {user}@{host}')


def main():
    user = input(f'Username [{default_user}]:') or default_user
    started_room_id, room_name = is_ctf_started(user)
    if not started_room_id:    # If no room running, proceed to selection and start
        login(user)
        get_ctf_list()
        ctf_id = input('Choose a CTF ID :')
        room_name = start_ctf(ctf_id)

    config = yes_or_no('Does this CTF need SSH and would like to configure it ? (Keys, password change...) [Y]/n?', default='Y')
    if config:
        config_ctf(room_name)


main()