Commit 0c5661b6 authored by Ondřej Borýsek's avatar Ondřej Borýsek
Browse files

Fix auto-login

parent ae3f54d3

pwndoc_api.py

+13 −6
Original line number Diff line number Diff line
@@ -23,14 +23,21 @@ class PwnDocAuth(AuthBase): 
            return r



        if session.cookies.get('token') is None:

            login()

            _login()

            r.prepare_cookies(session.cookies)  # Updating cookies on session does not automatically update cookies on prepared request.

            return r



        jwt_token = session.cookies.get('token', "").replace("JWT%20", "")

        jwt_data = jwt.decode(jwt_token, algorithms=["HS256"], options={"verify_signature": False})  # todo: why is the signature invalid?

        jwt_data = jwt.decode(jwt_token, algorithms=["HS256"], options={"verify_signature": False})

        exp_timestamp = jwt_data.get("exp")

        if current_timestamp() >= exp_timestamp:

            login()  # todo: maybe refresh token?



        remaining_seconds = exp_timestamp - current_timestamp()

        if remaining_seconds < 60:  # 60 seconds as a safe space

            try:

                _refresh_token()

            except AssertionError:

                _login()

            r.prepare_cookies(session.cookies)  # Updating cookies on session does not automatically update cookies on prepared request.



        return r
@@ -41,7 +48,7 @@ session.verify = not config.PWNDOC_DISABLE_HTTPS_VERIFICATION 
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)  # I'm writing one custom warning on config load, the default setting of urllib3 prints it on every request.





def login():

def _login():

    assert not config.PWNDOC_URL.startswith(config.PWNDOC_NO_CONNECTION_PLACEHOLDER), 'Tried to perform login NO-CONNECTION-TEST domain. Aborting.'

    data = {

        "username": config.PWNDOC_USERNAME,
@@ -53,7 +60,7 @@ def login(): 
    assert resp.status_code == 200, f"Login failed. Did you set PwnDoc username ({config.PWNDOC_USERNAME}) and password ({'*'*len(config.PWNDOC_PASSWORD)}) in ENV variables?"





def refresh_token():

def _refresh_token():

    resp = session.get(f"{config.PWNDOC_URL}/api/users/refreshtoken")

    assert resp.status_code == 200

pwndoc_db_init.py

+3 −3
Original line number Diff line number Diff line
@@ -10,7 +10,7 @@ from typing import Dict, Optional, List 
import config

from helpers.file_utils import relative_path



from pwndoc_api import session, login

from pwndoc_api import session





class InitialData:
@@ -32,7 +32,7 @@ class InitialData: 
            return



        self._add_pwndoc_user(first_user=True)

        # login()  # login happens automatically

        # login happens automatically



        # no dependencies

        self._upload_universal('_api_templates.json')
@@ -176,6 +176,7 @@ class InitialData: 
        # login is outside in PwnDocUpdate

        url = self._get_api_path_from_filename(CUSTOM_FIELDS_FILENAME)

        resp = session.get(url)

        assert resp.status_code == 200, f'Python Requests auth mechanism failed for PwnDoc failed? {resp.status_code}'

        pwndoc_datas = resp.json()["datas"]

        pwndoc_field = self.find_custom_field_by_label(pwndoc_datas, label)
@@ -186,6 +187,5 @@ class InitialData: 


class PwnDocUpdate:

    def add_new_fields_if_needed(self):

        login()

        _id = InitialData()

        _id.insert_vuln_custom_field_if_doesnt_exist('ignore_pluginoutput_checkbox')

tests/helpers_for_test.py

+3 −1
Original line number Diff line number Diff line 
import requests

import config

import pwndoc_api





def pwndoc_connectivity_test():

    ok = True

    try:

        pwndoc_api.login()

        resp = pwndoc_api.session.get(f'{config.PWNDOC_URL}/api/users/init')  # This is safe noop.

        assert resp.status_code == 200

    except requests.exceptions.ConnectionError:

        ok = False  # assert is not here to make the fail output shorter and more readable

    assert ok, "PwnDoc is not reachable. Is it running?"