get code from imap

app.py

@@ -6,6 +6,19 @@
 
 import steam.webauth as wa
 from os import environ
+import imaplib
+import email
+from time import sleep
+
+# get previously known message ids from imap server
+
+M = imaplib.IMAP4_SSL(environ['IMAP_HOST'])
+M.login(environ['IMAP_USER'], environ['IMAP_PASSWORD'])
+M.select()
+typ, data = M.search(None, '(FROM "noreply@steampowered.com" SUBJECT "Your Steam account: Access from new web or mobile device")')
+old_msg_ids = data[0].split()
+
+# perform steam login
 
 user = wa.WebAuth(environ['STEAM_USERNAME'])
 
@@ -17,7 +30,28 @@ except (wa.CaptchaRequired) as exp:
     print(user.captcha_url)
     user.login(password=environ['STEAM_PASSWORD'], captcha=input("Captcha: "))
 except wa.EmailCodeRequired:
-    user.login(email_code=input("Email Code: "))
+    while True:
+        typ, data = M.search(None, '(FROM "noreply@steampowered.com" SUBJECT "Your Steam account: Access from new web or mobile device")')
+        newest_msg_id = data[0].split()[-1]
+        
+        if newest_msg_id in old_msg_ids:
+            sleep(1)
+        else:
+            break
+
+    typ, data = M.fetch(newest_msg_id, '(RFC822)')
+    msg = email.message_from_bytes(data[0][1])
+    
+    for payload in msg.get_payload():
+        if (
+            payload.get_content_maintype() == 'text' and
+            payload.get_content_subtype() == 'plain'
+        ):
+            plaintext_lines = payload.get_payload(decode=True).decode().splitlines()
+            code = plaintext_lines[plaintext_lines.index('Login Code') + 1]
+            break
+    
+    user.login(email_code=code)
 except wa.TwoFactorCodeRequired:
     user.login(twofactor_code=input("2FA Code: "))
 
@@ -71,3 +105,8 @@ while True:
     html = r.json()['html']
     trs = BeautifulSoup(html, 'html.parser').find_all('tr')
     parse_trs(trs)
+
+# CLOSE
+
+M.close()
+M.logout()