From 66bd1b7dcc1abbb0c8e8580d964c053292ccd96e Mon Sep 17 00:00:00 2001
From: John Doe <miromaxxs@gmail.com>
Date: Mon, 7 Apr 2025 13:07:40 +0300
Subject: [PATCH] feat: add 2fa support, recreate session

---
 README.md           |  9 +++++++--
 Taskfile.yml        |  4 ++--
 auth.go             |  4 +++-
 internal/tg/auth.go | 10 +++++++---
 main.go             | 10 ++++++++++
 serve.go            |  2 +-
 6 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 71d4fa1..7a645ef 100644
--- a/README.md
+++ b/README.md
@@ -87,14 +87,19 @@ Before you can use the server, you need to connect to the Telegram API.
 
 1. Get the API ID and hash from [Telegram API](https://my.telegram.org/auth)
 2. Run the following command:
+   > __Note:__
+   > If you have 2FA enabled: add --password <2fa_password>
+
+   >  __Note:__
+   > If you want to override existing session: add --new
 
    ```bash
    telegram-mcp auth --app-id <your-api-id> --api-hash <your-api-hash> --phone <your-phone-number>
    ```
 
-   Enter the code you received from Telegram to connect to the API.
+   📩 Enter the code you received from Telegram to connect to the API.
 
-   The password may be required if you have two-factor authentication enabled.
+3. Done!
 
 ### Client Configuration
 
diff --git a/Taskfile.yml b/Taskfile.yml
index 7b79d85..be831bd 100644
--- a/Taskfile.yml
+++ b/Taskfile.yml
@@ -49,7 +49,7 @@ tasks:
   tag:
     desc: Create a new tag
     cmds:
-      - git tag -a v0.1.7
-      - git push origin v0.1.7
+      - git tag -a v0.1.8
+      - git push origin v0.1.8
 
 
diff --git a/auth.go b/auth.go
index dee7fcb..e85726e 100644
--- a/auth.go
+++ b/auth.go
@@ -13,6 +13,8 @@ import (
 
 func authCommand(_ context.Context, cmd *cli.Command) error {
 	phone := cmd.String("phone")
+	newSession := cmd.Bool("new")
+	pass := cmd.String("password")
 	appID := cmd.Root().Int("app-id")
 	apiHash := cmd.Root().String("api-hash")
 	sessionPath := cmd.Root().String("session")
@@ -24,7 +26,7 @@ func authCommand(_ context.Context, cmd *cli.Command) error {
 		Int64("app-id", appID).
 		Msg("Authenticate with Telegram")
 
-	err := tg.Auth(phone, appID, apiHash, sessionPath)
+	err := tg.Auth(phone, appID, apiHash, sessionPath, pass, newSession)
 	if err != nil {
 		log.Fatal().Err(err).Msg("Failed to authenticate with Telegram")
 	}
diff --git a/internal/tg/auth.go b/internal/tg/auth.go
index 80e0d71..b9ade09 100644
--- a/internal/tg/auth.go
+++ b/internal/tg/auth.go
@@ -14,7 +14,11 @@ import (
 	"github.com/rs/zerolog/log"
 )
 
-func Auth(phone string, appID int64, appHash string, sessionPath string) error {
+func Auth(phone string, appID int64, appHash string, sessionPath string, password string, newSession bool) error {
+	if newSession {
+		_ = os.Remove(sessionPath)
+	}
+
 	client := telegram.NewClient(int(appID), appHash, telegram.Options{
 		SessionStorage: &telegram.FileSessionStorage{
 			Path: sessionPath,
@@ -28,8 +32,8 @@ func Auth(phone string, appID int64, appHash string, sessionPath string) error {
 
 	if err := client.Run(context.Background(), func(ctx context.Context) error {
 		// Authenticate if needed
-		flow := auth.NewFlow(auth.Constant(phone, "", auth.CodeAuthenticatorFunc(func(_ context.Context, _ *tg.AuthSentCode) (string, error) {
-			fmt.Print("Enter code: ")
+		flow := auth.NewFlow(auth.Constant(phone, password, auth.CodeAuthenticatorFunc(func(_ context.Context, _ *tg.AuthSentCode) (string, error) {
+			fmt.Print("📩 Enter code: ")
 			code, err := bufio.NewReader(os.Stdin).ReadString('\n')
 			if err != nil {
 				return "", fmt.Errorf("read code: %w", err)
diff --git a/main.go b/main.go
index 788db2e..64b7442 100644
--- a/main.go
+++ b/main.go
@@ -77,6 +77,16 @@ func main() {
 						Required: true,
 						Aliases:  []string{"p"},
 					},
+					&cli.StringFlag{
+						Name:        "password",
+						Usage:       "Password for 2FA if exists",
+						HideDefault: true,
+					},
+					&cli.BoolFlag{
+						Name:        "new",
+						Usage:       "Remove old session and create new one",
+						HideDefault: true,
+					},
 				},
 				Action: authCommand,
 			},
diff --git a/serve.go b/serve.go
index bf94b10..be6d38f 100644
--- a/serve.go
+++ b/serve.go
@@ -48,7 +48,7 @@ func serve(ctx context.Context, cmd *cli.Command) error {
 
 		log.Info().RawJSON("answer", []byte(answer.Content[0].TextContent.Text)).Msg("Check GetDialogs: OK")
 
-		answer, err = client.GetHistory(tg.HistoryArguments{Name: "lalal", Offset: 5574})
+		answer, err = client.GetHistory(tg.HistoryArguments{Name: os.Getenv("TG_TEST_USERNAME")})
 		if err != nil {
 			return fmt.Errorf("get histore: %w", err)
 		}