[enhancement] Raise custom errors

[d-woosley]

I would love to see pwndocapi raise errors instead instead of just printing if in verbose. For example, I've been working on a script that will use the API to pull down a CSV but some users have MFA enabled. MFA support was added in my last pull, but I would like to be able to attempt the login, and if it fails, prompt for a TOTP and login again. However, I would like to only prompt for a TOTP if the failure is due to the lack of TOTP and not another general issue like a bad network connection or incorrect credentials.

Something like the following:

New Error Classes:

class PwnDocAPIException(Exception):
    pass

class MissingTOTP(PwnDocAPIException):
    def __init__(self, message, value):
        self.message = message
        self.value = value
        super().__init__(self.message)

class FailedLogin(PwnDocAPIException):
    def __init__(self, message, value):
        self.message = message
        self.value = value
        super().__init__(self.message)

class FailedConnection(PwnDocAPIException):
    def __init__(self, message, value):
        self.message = message
        self.value = value
        super().__init__(self.message)

Code Update Example:

from PwnDocAPIException import MissingTOTP

    def login(self, username, password, totp=""):
        r = self.session.post(
            self.target + "/api/users/token",
            json={"username": username, "password": password, "totpToken": totp},
            verify=False
        )
        if r.json()["status"] == "success":
            if self.verbose:
                print("[>] Successfully logged in.")
            self.user = {
                "token": r.json()["datas"]["token"],
                "refreshToken": r.json()["datas"]["refreshToken"]
            }
            self.session.cookies.set("token", "JWT%20" + self.user["token"])
            self.loggedin = True
        elif r.json()["status"] == "error":
            if self.verbose:
                print("[!] Login error. (%s)" % r.json()["datas"])
            self.loggedin = False

            # UPDATED CODE
            if r.json()["datas"] == "Example text about missing TOTP"
                raise MissingTOTP("authentication error", r.json()["datas"])
            if r.json()["datas"] == "Example text about failed login"
                raise FailedLogin("authentication error", r.json()["datas"])
            if r.json()["datas"] == "Example text about failed network connection"
                raise FailedConnection("authentication error", r.json()["datas"])

        return self.loggedin

Implementation in other scripts:

            # Start session
            session = pwndocapi.API(pwndoc_ip, 443, verbose=False)
           try:
                session.login(args.username, getpass("Password:"), totp="") 
           except MissingTOTP:
               session.login(args.username, getpass("Password:"), getpass("MFA TOTP:"))

I'm happy to work on this change and submit a pull, but I want to make sure this change is in-line with what is wanted for the project in general as this change could break other scripts. Of course, I can always fork pwndocapi and use my own copy, but I wanted to give back to the project if possible.

[p0dalirius]

Hi, Thank you for your suggestion I also think it would be useful to raise custom errors! Could you create a pull request for this? Best regards,

[d-woosley]

So sorry for the delayed response! I got a bit behind and had to switch projects. I will work on this when I can get to it next but it will likely be a little while (a few months).