spider/internal/telegram/sessionimport/importer.go

package sessionimport

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"

	"github.com/gotd/td/session"
)

// PyrogramMeta is the parsed content of <phone>.json inside a 协议号 folder.
// Unknown fields are ignored.
type PyrogramMeta struct {
	Phone          string `json:"phone"`
	APIID          int    `json:"api_id"`
	APIHash        string `json:"api_hash"`
	AppID          int    `json:"app_id"`   // some generators use app_id/app_hash instead
	AppHash        string `json:"app_hash"`
	Device         string `json:"device"`
	AppVersion     string `json:"app_version"`
	SDK            string `json:"sdk"`
	LangPack       string `json:"lang_pack"`
	LangCode       string `json:"lang_code"`
	SystemLangCode string `json:"system_lang_code"`
	FirstName      string `json:"first_name"`
	LastName       string `json:"last_name"`
	Username       string `json:"username"`
	TwoFA          string `json:"twoFA"`
}

// Normalize resolves API credential aliasing. 协议号 JSON files use api_id/api_hash
// or app_id/app_hash interchangeably; we coalesce them.
func (m *PyrogramMeta) Normalize() {
	if m.APIID == 0 && m.AppID != 0 {
		m.APIID = m.AppID
	}
	if m.APIHash == "" && m.AppHash != "" {
		m.APIHash = m.AppHash
	}
}

// Result is the outcome of importing one protocol-number folder.
type Result struct {
	Meta           PyrogramMeta
	TwoFAPlaintext string      // from 2fa_password.txt if Meta.TwoFA empty
	SessionData    *session.Data
	UsedSource     string      // "pyrogram" or "tdata"
	SessionFile    string      // sessions/<phone>.json
	OriginDir      string      // sessions/<phone>.origin/
	PyrogramErr    string      // empty on success
	TdataErr       string      // empty on success or if tdata absent
}

// StagedFile is one file extracted from a multipart upload, already written to a
// temp directory. RelPath preserves the browser's webkitRelativePath minus the
// top-level folder name (so e.g. "13252753163/tdata/key_datas" becomes
// "tdata/key_datas").
type StagedFile struct {
	RelPath string
	AbsPath string
}

// Import parses staged files from a 协议号 folder, converts the session, and
// materializes sessions/<phone>.json + origin/ backup under sessionsDir.
// The caller (HTTP handler) persists the returned metadata to the DB.
//
// Precondition: all StagedFile.AbsPath must exist on disk. stagedFiles is the
// output of the handler's multipart parser.
func Import(ctx context.Context, sessionsDir string, stagedFiles []StagedFile) (*Result, error) {
	if len(stagedFiles) == 0 {
		return nil, errors.New("no files uploaded")
	}

	var metaPath, sessionPath, twoFAPath, tdataDir string
	for _, f := range stagedFiles {
		switch {
		case filepath.Ext(f.RelPath) == ".json" && !hasSep(f.RelPath):
			metaPath = f.AbsPath
		case filepath.Ext(f.RelPath) == ".session":
			sessionPath = f.AbsPath
		case filepath.Base(f.RelPath) == "2fa_password.txt":
			twoFAPath = f.AbsPath
		case hasTdataPrefix(f.RelPath):
			// Remember the parent dir containing tdata/
			dir := filepath.Dir(f.AbsPath)
			// Walk up until we find a folder literally named "tdata".
			for dir != "" && filepath.Base(dir) != "tdata" {
				dir = filepath.Dir(dir)
			}
			if dir != "" {
				tdataDir = dir
			}
		}
	}

	if metaPath == "" {
		return nil, errors.New("missing <phone>.json in upload")
	}
	if sessionPath == "" && tdataDir == "" {
		return nil, errors.New("upload has neither .session nor tdata/")
	}

	meta, err := readMeta(metaPath)
	if err != nil {
		return nil, err
	}
	meta.Normalize()
	if meta.Phone == "" {
		return nil, errors.New("<phone>.json missing phone field")
	}
	if meta.APIID == 0 || meta.APIHash == "" {
		return nil, errors.New("<phone>.json missing api_id / api_hash")
	}

	res := &Result{Meta: meta}

	// Prefer 2FA from JSON; fall back to 2fa_password.txt.
	if meta.TwoFA == "" && twoFAPath != "" {
		b, err := os.ReadFile(twoFAPath)
		if err == nil {
			res.TwoFAPlaintext = sanitizeTwoFA(string(b))
		}
	} else {
		res.TwoFAPlaintext = meta.TwoFA
	}

	// Try Pyrogram first.
	if sessionPath != "" {
		data, err := ConvertPyrogramSession(sessionPath)
		if err == nil {
			res.SessionData = data
			res.UsedSource = "pyrogram"
		} else {
			res.PyrogramErr = err.Error()
		}
	}

	// Fallback to tdata.
	if res.SessionData == nil && tdataDir != "" {
		data, err := ConvertTdataSession(tdataDir, res.TwoFAPlaintext)
		if err == nil {
			res.SessionData = data
			res.UsedSource = "tdata"
		} else {
			res.TdataErr = err.Error()
		}
	}

	if res.SessionData == nil {
		return res, fmt.Errorf("session conversion failed (pyrogram: %s; tdata: %s)",
			nonempty(res.PyrogramErr, "n/a"), nonempty(res.TdataErr, "n/a"))
	}

	// Materialize session JSON.
	if err := os.MkdirAll(sessionsDir, 0o755); err != nil {
		return res, fmt.Errorf("create sessions dir: %w", err)
	}
	res.SessionFile = filepath.Join(sessionsDir, meta.Phone+".json")
	storage := &session.FileStorage{Path: res.SessionFile}
	loader := session.Loader{Storage: storage}
	if err := loader.Save(ctx, res.SessionData); err != nil {
		return res, fmt.Errorf("save gotd session: %w", err)
	}

	// Materialize origin backup (best-effort; non-fatal).
	res.OriginDir = filepath.Join(sessionsDir, meta.Phone+".origin")
	_ = os.RemoveAll(res.OriginDir)
	if err := os.MkdirAll(res.OriginDir, 0o755); err == nil {
		for _, f := range stagedFiles {
			dst := filepath.Join(res.OriginDir, f.RelPath)
			if err := os.MkdirAll(filepath.Dir(dst), 0o755); err == nil {
				_ = copyFile(f.AbsPath, dst)
			}
		}
	}

	return res, nil
}

func readMeta(path string) (PyrogramMeta, error) {
	var m PyrogramMeta
	b, err := os.ReadFile(path)
	if err != nil {
		return m, fmt.Errorf("read meta json: %w", err)
	}
	if err := json.Unmarshal(b, &m); err != nil {
		return m, fmt.Errorf("parse meta json: %w", err)
	}
	return m, nil
}

func sanitizeTwoFA(s string) string {
	// Trim CR/LF/space — 2fa_password.txt often has a trailing newline.
	for len(s) > 0 && (s[len(s)-1] == '\n' || s[len(s)-1] == '\r' || s[len(s)-1] == ' ') {
		s = s[:len(s)-1]
	}
	return s
}

func hasSep(p string) bool {
	for i := 0; i < len(p); i++ {
		if p[i] == '/' || p[i] == '\\' {
			return true
		}
	}
	return false
}

func hasTdataPrefix(p string) bool {
	// "tdata/..." or "tdata\..."
	if len(p) < 6 {
		return false
	}
	return (p[:5] == "tdata" && (p[5] == '/' || p[5] == '\\'))
}

func nonempty(s, alt string) string {
	if s == "" {
		return alt
	}
	return s
}

func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, in)
	return err
}