request.py

"""
This module abstracts the requests.sessions package to
detect DCAR SSO.
"""

import sys

from time import sleep
from requests import session, Response
from requests.exceptions import ProxyError

import bs4 as bs # type: ignore
import pyotp # type: ignore

from dcar.utils import validate_argument

class LoginError(Exception):
    """DCAR Login Exception"""

class Session:
    """Wrapper class for requests sessions which detects if SSO needs
       to be called prior to calling actual request"""

    def __init__(self, username: str, password: str, totp_seed: str):
        validate_argument("username", username, str)
        validate_argument("password", password, str)
        validate_argument("totp_seed", totp_seed, str)

        self.session = session()
        self.username = username
        self.password = password
        self.totp_seed = totp_seed.replace(" ", "")

    def login(self, dcar_url: str = "https://dcar.dso.mil"):
        """
        Method which will authenicate with the DCAR website.
        """
        validate_argument("dcar_url", dcar_url, str)

        totp = pyotp.TOTP(self.totp_seed)
        try:
            # jump through login page
            login_request = self.session.get(dcar_url, verify=False)
            login_page = bs.BeautifulSoup(login_request.content, 'html.parser')
            login_form_action_url = login_page.find('form').get('action')

            # submit time-based one-time code
            totp_request = self.session.post(login_form_action_url,
                                             data={'username': self.username, 'password': self.password, 'credentialId': ''}, verify=False)
            totp_page = bs.BeautifulSoup(totp_request.content, 'html.parser')
            totp_form_action_url = totp_page.find('form').get('action')

            # Submit agreement to user agreement
            user_agreement_request = self.session.post(totp_form_action_url, data={'otp': totp.now(), 'login': 'MFA Log In'}, verify=False)
            user_agreement_page = bs.BeautifulSoup(user_agreement_request.content, 'html.parser')
            user_agreement_url = user_agreement_page.find('form').get('action')

            # finally authenicated and accessing the site
            finished_request = self.session.post(user_agreement_url, data={'accept': 'Accept'}, verify=False)
            finished_page = bs.BeautifulSoup(finished_request.content, 'html.parser')

            if finished_page.find('input', id='password-new'):
                raise LoginError("Need to update password for Iron Bank")

        # AttributeError exception occurs when beautiful soup fails to find a id.
        # Typically this will be the 'kc-form-login' or 'kc-totp-login-form' ids.
        except AttributeError:
            raise LoginError("Error logging into Iron Bank, login form likey changed")

    def get(self, path: str = '', stream: bool = False, verify: bool = False, max_retries: int = 10,
            retry_sleep: int = 10, dcar_url: str = "https://ironbank.dso.mil",
            sso_url: str = "https://login.dso.mil") -> Response:
        """
        Wrapper method for the requests.get() method. This method
        will check if the request was redirected to DCAR SSO. If
        redirected, then login before returning the authenicated
        response.
        """
        validate_argument("path", path, str)
        validate_argument("stream", stream, bool)
        validate_argument("verify", verify, bool)
        validate_argument("max_retries", max_retries, int)
        validate_argument("retry_sleep", retry_sleep, int)
        validate_argument("dcar_url", dcar_url, str)
        validate_argument("sso_url", sso_url, str)

        for _ in range(max_retries):
            response = self.session.get(dcar_url + path, stream=stream, verify=verify, timeout=20)
            if not sso_url in response.url:
                return response
            self.login(dcar_url)
            sleep(retry_sleep)
        raise LoginError("Unable to authenicate with Iron Bank")