diff --git a/README.md b/README.md index f7f76a9..fb4d9bc 100644 --- a/README.md +++ b/README.md @@ -30,7 +30,7 @@ - **Search** — Search chats, messages, and global Telegram directory - **Contacts** — Contact list with online status indicators - **Group Info** — Member list, admin roles, group description -- **Authentication** — Phone/SMS code, 2FA password (QR code login not yet implemented) +- **Authentication** — Phone/SMS code and 2FA password, plus QR login for `telegram-mcp` - **First-Run Wizard** — Prompts for API credentials and saves config automatically - **Notifications** — Desktop notifications via `notify-send` / `osascript` - **Responsive Layout** — Dual-panel (wide) or single-panel (narrow terminals) @@ -232,9 +232,15 @@ The repo also ships `telegram-mcp`, an [MCP](https://modelcontextprotocol.io) se The MCP server uses a separate session (`session-mcp.json`), so log in once even if the TUI is already logged in: ```bash -bin/telegram-mcp login # phone → code → 2FA, writes ~/.local/share/tele-tui/session-mcp.json +bin/telegram-mcp login # phone → code → 2FA +bin/telegram-mcp login --qr # scan in Telegram → Settings → Devices ``` +QR tokens refresh automatically until the login is accepted or cancelled. If +the account has two-step verification enabled, the password is read without +echoing it to the terminal. Both login modes write +`~/.local/share/tele-tui/session-mcp.json` by default. + ### Client configuration Register the server in your MCP client, e.g.: diff --git a/cmd/telegram-mcp/main.go b/cmd/telegram-mcp/main.go index e8a9e65..211e568 100644 --- a/cmd/telegram-mcp/main.go +++ b/cmd/telegram-mcp/main.go @@ -3,15 +3,18 @@ // Subcommands: // // serve (default) run the MCP server on stdin/stdout -// login interactive login, writes the MCP session file -// (~/.local/share/tele-tui/session-mcp.json by default) +// login interactive phone login, writes the MCP session file +// login --qr QR login via an already authorized Telegram app +// (~/.local/share/tele-tui/session-mcp.json by default) package main import ( "bufio" "context" "errors" + "flag" "fmt" + "io" "log" "os" "os/signal" @@ -20,9 +23,12 @@ import ( "syscall" "time" + "golang.org/x/term" + "github.com/imtaqin/telegram-cli/internal/config" "github.com/imtaqin/telegram-cli/internal/mcpserver" "github.com/imtaqin/telegram-cli/internal/telegram" + "github.com/imtaqin/telegram-cli/internal/ui/widgets" ) const loginHint = "session not authorized, run 'telegram-mcp login' first" @@ -32,9 +38,15 @@ func main() { log.SetOutput(os.Stderr) log.SetPrefix("telegram-mcp: ") - cmd := "serve" - if len(os.Args) > 1 { - cmd = os.Args[1] + opts, err := parseCommand(os.Args[1:]) + if errors.Is(err, flag.ErrHelp) { + printUsage(os.Stdout) + return + } + if err != nil { + fmt.Fprintf(os.Stderr, "telegram-mcp: %v\n", err) + printUsage(os.Stderr) + os.Exit(2) } cfg, err := config.Load() @@ -55,15 +67,53 @@ func main() { cfg.Storage.SessionFile = strings.TrimSuffix(cfg.Storage.SessionFile, ".json") + "-mcp.json" } - switch cmd { + switch opts.command { case "login": - runLogin(cfg) + if opts.qr { + runQRLogin(cfg) + } else { + runLogin(cfg) + } case "serve": runServe(cfg) + } +} + +type commandOptions struct { + command string + qr bool +} + +func parseCommand(args []string) (commandOptions, error) { + opts := commandOptions{command: "serve"} + if len(args) > 0 && !strings.HasPrefix(args[0], "-") { + opts.command = args[0] + args = args[1:] + } + + switch opts.command { + case "login", "serve": default: - fmt.Fprintf(os.Stderr, "usage: telegram-mcp [login|serve]\n") - os.Exit(2) + return commandOptions{}, fmt.Errorf("unknown command %q", opts.command) + } + + fs := flag.NewFlagSet("telegram-mcp "+opts.command, flag.ContinueOnError) + fs.SetOutput(io.Discard) + fs.BoolVar(&opts.qr, "qr", false, "log in by scanning a QR code") + if err := fs.Parse(args); err != nil { + return commandOptions{}, err + } + if fs.NArg() != 0 { + return commandOptions{}, fmt.Errorf("unexpected argument %q", fs.Arg(0)) } + if opts.qr && opts.command != "login" { + return commandOptions{}, errors.New("--qr is only valid with the login command") + } + return opts, nil +} + +func printUsage(w io.Writer) { + fmt.Fprintln(w, "usage: telegram-mcp [login [--qr]|serve]") } // runLogin performs interactive authentication in the terminal and @@ -140,6 +190,47 @@ func runLogin(cfg *config.Config) { log.Fatalf("login succeeded but GetMe failed: %v", err) } + printLoggedIn(me) +} + +func runQRLogin(cfg *config.Config) { + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + + me, err := telegram.LoginWithQR(ctx, cfg, telegram.QRLoginOptions{ + ShowQRCode: func(_ context.Context, token telegram.QRLoginToken) error { + fmt.Fprint(os.Stderr, "\x1b[2J\x1b[H") + fmt.Fprintln(os.Stderr, "Telegram QR Login") + fmt.Fprintln(os.Stderr) + fmt.Fprintln(os.Stderr, "On your logged-in Telegram phone:") + fmt.Fprintln(os.Stderr, "Settings -> Devices -> Link Desktop Device") + fmt.Fprintln(os.Stderr) + fmt.Fprintln(os.Stderr, widgets.RenderQRCode(token.URL, 256)) + fmt.Fprintln(os.Stderr) + fmt.Fprintf(os.Stderr, "QR expires at %s and will refresh automatically.\n", token.ExpiresAt.Format("15:04:05")) + return nil + }, + PasswordPrompt: func(_ context.Context, retry bool) ([]byte, error) { + if retry { + fmt.Fprintln(os.Stderr, "The password was empty or invalid. Try again.") + } + fmt.Fprint(os.Stderr, "Enter your Telegram 2FA password (input is hidden): ") + password, err := term.ReadPassword(int(os.Stdin.Fd())) + fmt.Fprintln(os.Stderr) + return password, err + }, + }) + if err != nil { + if errors.Is(err, context.Canceled) { + return + } + log.Fatalf("QR login failed: %v", err) + } + + printLoggedIn(me) +} + +func printLoggedIn(me *telegram.User) { name := strings.TrimSpace(me.FirstName + " " + me.LastName) if me.Username != "" { fmt.Fprintf(os.Stderr, "Logged in as %s (@%s, id %d)\n", name, me.Username, me.ID) diff --git a/cmd/telegram-mcp/main_test.go b/cmd/telegram-mcp/main_test.go new file mode 100644 index 0000000..ffecb4a --- /dev/null +++ b/cmd/telegram-mcp/main_test.go @@ -0,0 +1,86 @@ +package main + +import ( + "bytes" + "os" + "os/exec" + "strings" + "testing" +) + +const helpTestProcessEnv = "TELEGRAM_MCP_HELP_TEST_PROCESS" + +func TestParseCommand(t *testing.T) { + tests := []struct { + name string + args []string + command string + qr bool + wantErr bool + }{ + {name: "default serve", command: "serve"}, + {name: "phone login", args: []string{"login"}, command: "login"}, + {name: "QR login", args: []string{"login", "--qr"}, command: "login", qr: true}, + {name: "serve rejects QR flag", args: []string{"serve", "--qr"}, wantErr: true}, + {name: "default serve rejects QR flag", args: []string{"--qr"}, wantErr: true}, + {name: "unknown command", args: []string{"unknown"}, wantErr: true}, + {name: "unexpected argument", args: []string{"login", "extra"}, wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := parseCommand(tt.args) + if tt.wantErr { + if err == nil { + t.Fatalf("expected an error, got %+v", got) + } + return + } + if err != nil { + t.Fatalf("parse command: %v", err) + } + if got.command != tt.command || got.qr != tt.qr { + t.Fatalf("expected command=%q qr=%t, got command=%q qr=%t", tt.command, tt.qr, got.command, got.qr) + } + }) + } +} + +func TestMainHelpExitsSuccessfully(t *testing.T) { + if args, ok := os.LookupEnv(helpTestProcessEnv); ok { + os.Args = append([]string{"telegram-mcp"}, strings.Fields(args)...) + main() + return + } + + tests := []struct { + name string + args []string + }{ + {name: "long flag", args: []string{"--help"}}, + {name: "short flag", args: []string{"-h"}}, + {name: "login help", args: []string{"login", "--help"}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cmd := exec.Command(os.Args[0], "-test.run=^TestMainHelpExitsSuccessfully$") + cmd.Env = append(os.Environ(), helpTestProcessEnv+"="+strings.Join(tt.args, " ")) + + var stdout bytes.Buffer + var stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + if err := cmd.Run(); err != nil { + t.Fatalf("telegram-mcp %s failed: %v\nstderr:\n%s", strings.Join(tt.args, " "), err, stderr.String()) + } + if !strings.Contains(stdout.String(), "usage: telegram-mcp [login [--qr]|serve]") { + t.Fatalf("expected usage on stdout, got %q", stdout.String()) + } + if stderr.Len() != 0 { + t.Fatalf("expected empty stderr, got %q", stderr.String()) + } + }) + } +} diff --git a/go.mod b/go.mod index aede45c..440aed0 100644 --- a/go.mod +++ b/go.mod @@ -11,6 +11,7 @@ require ( github.com/pelletier/go-toml/v2 v2.2.3 github.com/skip2/go-qrcode v0.0.0-20200617195104-da1b6568686e golang.org/x/image v0.38.0 + golang.org/x/term v0.45.0 ) require ( @@ -75,7 +76,6 @@ require ( golang.org/x/oauth2 v0.35.0 // indirect golang.org/x/sync v0.22.0 // indirect golang.org/x/sys v0.47.0 // indirect - golang.org/x/term v0.45.0 // indirect golang.org/x/text v0.40.0 // indirect golang.org/x/time v0.15.0 // indirect golang.org/x/tools v0.48.0 // indirect diff --git a/internal/telegram/qr.go b/internal/telegram/qr.go new file mode 100644 index 0000000..dc071c1 --- /dev/null +++ b/internal/telegram/qr.go @@ -0,0 +1,175 @@ +package telegram + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "time" + + "github.com/gotd/td/session" + gotd "github.com/gotd/td/telegram" + gotdauth "github.com/gotd/td/telegram/auth" + "github.com/gotd/td/telegram/auth/qrlogin" + "github.com/gotd/td/tg" + "github.com/gotd/td/tgerr" + + "github.com/imtaqin/telegram-cli/internal/config" +) + +// ErrQRPasswordPromptRequired is returned when QR login reaches 2FA but no +// password prompt was configured. +var ErrQRPasswordPromptRequired = errors.New("QR login requires a 2FA password prompt") + +// QRLoginToken is a short-lived Telegram login token to render as a QR code. +type QRLoginToken struct { + URL string + ExpiresAt time.Time +} + +// QRLoginOptions supplies the interactive parts of QR authentication. +type QRLoginOptions struct { + // ShowQRCode is called whenever Telegram issues or refreshes a QR token. + ShowQRCode func(context.Context, QRLoginToken) error + + // PasswordPrompt is called if the account requires 2FA. retry is true + // after an empty or invalid password. The returned byte slice is consumed + // and wiped before the function returns. + PasswordPrompt func(ctx context.Context, retry bool) ([]byte, error) +} + +// LoginWithQR authorizes cfg.Storage.SessionFile by scanning a QR code in an +// already authorized Telegram app. Expired QR tokens are refreshed by gotd. +func LoginWithQR(ctx context.Context, cfg *config.Config, opts QRLoginOptions) (*User, error) { + if opts.ShowQRCode == nil { + return nil, errors.New("QR login requires a QR code callback") + } + if err := os.MkdirAll(filepath.Dir(cfg.Storage.SessionFile), 0o755); err != nil { + return nil, fmt.Errorf("create session directory: %w", err) + } + + dispatcher := tg.NewUpdateDispatcher() + loggedIn := qrlogin.OnLoginToken(dispatcher) + client := gotd.NewClient(int(cfg.Telegram.APIID), cfg.Telegram.APIHash, gotd.Options{ + SessionStorage: &session.FileStorage{Path: cfg.Storage.SessionFile}, + UpdateHandler: dispatcher, + Device: gotd.DeviceConfig{ + DeviceModel: "Telegram CLI", + SystemVersion: "1.0.0", + AppVersion: "0.1.0", + SystemLangCode: "en", + LangCode: "en", + }, + }) + + var result *User + err := client.Run(ctx, func(ctx context.Context) error { + status, err := client.Auth().Status(ctx) + if err != nil { + return fmt.Errorf("check authorization status: %w", err) + } + if status.Authorized { + if status.User == nil { + return errors.New("authorized Telegram session has no user") + } + result = userFromTG(status.User) + return nil + } + + err = finishQRAuthentication( + ctx, + func(ctx context.Context) error { + _, err := client.QR().Auth(ctx, loggedIn, func(ctx context.Context, token qrlogin.Token) error { + return opts.ShowQRCode(ctx, QRLoginToken{ + URL: token.URL(), + ExpiresAt: token.Expires(), + }) + }) + return err + }, + func(ctx context.Context) error { + return completeQRPassword(ctx, client.Auth(), opts.PasswordPrompt) + }, + ) + if err != nil { + return fmt.Errorf("QR authentication: %w", err) + } + + status, err = client.Auth().Status(ctx) + if err != nil { + return fmt.Errorf("verify authorization status: %w", err) + } + if !status.Authorized || status.User == nil { + return errors.New("Telegram did not authorize the QR session") + } + result = userFromTG(status.User) + return nil + }) + if err != nil { + return nil, err + } + if result == nil { + return nil, errors.New("QR login completed without an authorized user") + } + return result, nil +} + +func finishQRAuthentication( + ctx context.Context, + qrAuth func(context.Context) error, + passwordAuth func(context.Context) error, +) error { + err := qrAuth(ctx) + if err == nil { + return nil + } + if !tgerr.Is(err, "SESSION_PASSWORD_NEEDED") { + return err + } + return passwordAuth(ctx) +} + +type qrPasswordClient interface { + PasswordWith(context.Context, gotdauth.PasswordHashFunc) (*tg.AuthAuthorization, error) +} + +func completeQRPassword( + ctx context.Context, + client qrPasswordClient, + prompt func(context.Context, bool) ([]byte, error), +) error { + if prompt == nil { + return ErrQRPasswordPromptRequired + } + + retry := false + for { + secret, err := prompt(ctx, retry) + if err != nil { + wipeBytes(secret) + return fmt.Errorf("read 2FA password: %w", err) + } + if len(secret) == 0 { + retry = true + continue + } + + _, err = client.PasswordWith(ctx, gotdauth.PasswordHashFor(secret)) + wipeBytes(secret) + if err == nil { + return nil + } + if errors.Is(err, gotdauth.ErrPasswordInvalid) { + retry = true + continue + } + return fmt.Errorf("verify 2FA password: %w", err) + } +} + +func wipeBytes(value []byte) { + for i := range value { + value[i] = 0 + } +} diff --git a/internal/telegram/qr_test.go b/internal/telegram/qr_test.go new file mode 100644 index 0000000..beebb07 --- /dev/null +++ b/internal/telegram/qr_test.go @@ -0,0 +1,121 @@ +package telegram + +import ( + "context" + "errors" + "fmt" + "testing" + + gotdauth "github.com/gotd/td/telegram/auth" + "github.com/gotd/td/tg" + "github.com/gotd/td/tgerr" +) + +func TestFinishQRAuthenticationFallsBackTo2FA(t *testing.T) { + passwordCalled := false + err := finishQRAuthentication( + context.Background(), + func(context.Context) error { + return fmt.Errorf("export login token: %w", tgerr.New(401, "SESSION_PASSWORD_NEEDED")) + }, + func(context.Context) error { + passwordCalled = true + return nil + }, + ) + + if err != nil { + t.Fatalf("finish QR authentication: %v", err) + } + if !passwordCalled { + t.Fatal("expected the 2FA password flow to run") + } +} + +func TestFinishQRAuthenticationReturnsOtherErrors(t *testing.T) { + want := errors.New("QR export failed") + passwordCalled := false + err := finishQRAuthentication( + context.Background(), + func(context.Context) error { return want }, + func(context.Context) error { + passwordCalled = true + return nil + }, + ) + + if !errors.Is(err, want) { + t.Fatalf("expected %v, got %v", want, err) + } + if passwordCalled { + t.Fatal("did not expect the 2FA password flow to run") + } +} + +func TestCompleteQRPasswordRetriesAndWipesSecrets(t *testing.T) { + first := []byte("wrong password") + second := []byte("correct password") + promptCalls := 0 + client := &fakeQRPasswordClient{errors: []error{gotdauth.ErrPasswordInvalid, nil}} + + err := completeQRPassword(context.Background(), client, func(_ context.Context, retry bool) ([]byte, error) { + if retry != (promptCalls > 0) { + t.Fatalf("unexpected retry value %t on prompt call %d", retry, promptCalls+1) + } + promptCalls++ + if promptCalls == 1 { + return first, nil + } + return second, nil + }) + + if err != nil { + t.Fatalf("complete QR password: %v", err) + } + if promptCalls != 2 { + t.Fatalf("expected 2 password prompts, got %d", promptCalls) + } + assertZeroed(t, first) + assertZeroed(t, second) +} + +func TestCompleteQRPasswordWipesSecretWhenPromptFails(t *testing.T) { + secret := []byte("partial password") + wantErr := errors.New("interrupted read") + client := &fakeQRPasswordClient{} + + err := completeQRPassword(context.Background(), client, func(context.Context, bool) ([]byte, error) { + return secret, wantErr + }) + + if !errors.Is(err, wantErr) { + t.Fatalf("expected prompt error %v, got %v", wantErr, err) + } + if client.calls != 0 { + t.Fatalf("expected no password verification calls, got %d", client.calls) + } + assertZeroed(t, secret) +} + +func assertZeroed(t *testing.T, secret []byte) { + t.Helper() + for i, b := range secret { + if b != 0 { + t.Fatalf("secret byte %d was not wiped", i) + } + } +} + +type fakeQRPasswordClient struct { + errors []error + calls int +} + +func (f *fakeQRPasswordClient) PasswordWith(context.Context, gotdauth.PasswordHashFunc) (*tg.AuthAuthorization, error) { + err := f.errors[f.calls] + f.calls++ + if err != nil { + return nil, err + } + return &tg.AuthAuthorization{}, nil +}