first commit

cmd/root.go

@@ -0,0 +1,83 @@
+package cmd
+
+import (
+	"fmt"
+	"log"
+	"net/http"
+	"os"
+	"otpm/database"
+	"otpm/handlers"
+	"otpm/utils"
+
+	"github.com/jmoiron/sqlx"
+	"github.com/julienschmidt/httprouter"
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+)
+
+var rootCmd = &cobra.Command{
+	Use:   "otpm",
+	Short: "otp backend for microapp on wechat",
+	Run: func(cmd *cobra.Command, args []string) {
+		startApp()
+	},
+}
+
+func Execute() {
+	if err := rootCmd.Execute(); err != nil {
+		fmt.Println(err)
+		os.Exit(1)
+	}
+}
+
+func init() {
+	cobra.OnInitialize(initConfig)
+
+	rootCmd.PersistentFlags().StringP("config", "c", "", "config file (default is $HOME/config.yaml)")
+	rootCmd.PersistentFlags().StringP("driver", "d", "sqlite3", "database driver (sqlite3, postgres, mysql)")
+	rootCmd.PersistentFlags().StringP("dsn", "s", "", "database connection string")
+	rootCmd.PersistentFlags().StringP("port", "p", "8080", "port to listen on")
+
+	viper.BindPFlag("database.driver", rootCmd.PersistentFlags().Lookup("driver"))
+	viper.BindPFlag("database.dsn", rootCmd.PersistentFlags().Lookup("dsn"))
+	viper.BindPFlag("port", rootCmd.PersistentFlags().Lookup("port"))
+}
+
+func initConfig() {
+	if cfgFile := viper.GetString("config"); cfgFile != "" {
+		viper.SetConfigFile(cfgFile)
+	} else {
+		viper.AddConfigPath(".")
+		viper.SetConfigName("config")
+		viper.SetConfigType("yaml")
+	}
+
+	if err := viper.ReadInConfig(); err != nil {
+		log.Fatalf("Error reading config file: %v", err)
+	}
+}
+
+func initApp(db *sqlx.DB) {
+	if err := database.MigrateDB(db); err != nil {
+		log.Fatalf("Error migrating the database: %v", err)
+	}
+}
+
+func startApp() {
+	port := viper.GetInt("port")
+	db, err := database.InitDB()
+	if err != nil {
+		log.Fatalf("Error connecting to the database: %v", err)
+	}
+	defer db.Close()
+	initApp(db)
+	handler := &handlers.Handler{DB: db}
+
+	router := httprouter.New()
+	router.POST("/login", utils.AdaptHandler(handler.Login))
+	router.POST("/set", utils.AdaptHandler(handler.UpdateOrCreateOtp))
+	router.GET("/get", utils.AdaptHandler(handler.GetOtp))
+
+	log.Println("Starting server on :8080")
+	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", port), router))
+}