From 26447d19314197d0b11f3e1e0f41f11db39c8be4 Mon Sep 17 00:00:00 2001
From: JonnyWong16 <9099342+JonnyWong16@users.noreply.github.com>
Date: Sat, 17 Feb 2024 14:34:43 -0800
Subject: [PATCH] Add exception for two-factor required (#1357)

* Add exception for two-factor required

* Update tools/plex-gettoken.py with 2FA exception
---
 plexapi/exceptions.py  | 5 +++++
 plexapi/myplex.py      | 4 +++-
 tools/plex-gettoken.py | 9 +++++++--
 3 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/plexapi/exceptions.py b/plexapi/exceptions.py
index c269c38e..182feb13 100644
--- a/plexapi/exceptions.py
+++ b/plexapi/exceptions.py
@@ -29,3 +29,8 @@ class Unsupported(PlexApiException):
 class Unauthorized(BadRequest):
     """ Invalid username/password or token. """
     pass
+
+
+class TwoFactorRequired(Unauthorized):
+    """ Two factor authentication required. """
+    pass
diff --git a/plexapi/myplex.py b/plexapi/myplex.py
index 461cfecb..8d697924 100644
--- a/plexapi/myplex.py
+++ b/plexapi/myplex.py
@@ -12,7 +12,7 @@ from plexapi import (BASE_HEADERS, CONFIG, TIMEOUT, X_PLEX_ENABLE_FAST_CONNECT,
                      log, logfilter, utils)
 from plexapi.base import PlexObject
 from plexapi.client import PlexClient
-from plexapi.exceptions import BadRequest, NotFound, Unauthorized
+from plexapi.exceptions import BadRequest, NotFound, Unauthorized, TwoFactorRequired
 from plexapi.library import LibrarySection
 from plexapi.server import PlexServer
 from plexapi.sonos import PlexSonosClient
@@ -237,6 +237,8 @@ class MyPlexAccount(PlexObject):
             errtext = response.text.replace('\n', ' ')
             message = f'({response.status_code}) {codename}; {response.url} {errtext}'
             if response.status_code == 401:
+                if "verification code" in response.text:
+                    raise TwoFactorRequired(message)
                 raise Unauthorized(message)
             elif response.status_code == 404:
                 raise NotFound(message)
diff --git a/tools/plex-gettoken.py b/tools/plex-gettoken.py
index 48b60a8b..03c32ae1 100755
--- a/tools/plex-gettoken.py
+++ b/tools/plex-gettoken.py
@@ -4,11 +4,16 @@
 Plex-GetToken is a simple method to retrieve a Plex account token.
 """
 from getpass import getpass
+from plexapi.exceptions import TwoFactorRequired
 from plexapi.myplex import MyPlexAccount
 
 username = input("Plex username: ")
 password = getpass("Plex password: ")
-code = input("Plex 2FA code (leave blank for none): ")
 
-account = MyPlexAccount(username, password, code=code)
+try:
+    account = MyPlexAccount(username, password)
+except TwoFactorRequired:
+    code = input("Plex 2FA code: ")
+    account = MyPlexAccount(username, password, code=code)
+
 print(account.authenticationToken)