arimelody.me/main.go

package main

import (
	"errors"
	"fmt"
	stdLog "log"
	"math"
	"math/rand"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"arimelody-web/admin"
	"arimelody-web/api"
	"arimelody-web/colour"
	"arimelody-web/controller"
	"arimelody-web/model"
	"arimelody-web/templates"
	"arimelody-web/view"
  "arimelody-web/log"

	"github.com/jmoiron/sqlx"
	_ "github.com/lib/pq"
	"golang.org/x/crypto/bcrypt"
)

// used for database migrations
const DB_VERSION = 1

const DEFAULT_PORT int64 = 8080

func main() {
    fmt.Printf("made with <3 by ari melody\n\n")

    app := model.AppState{
        Config: controller.GetConfig(),
    }

    // initialise database connection
    if app.Config.DB.Host == "" {
        fmt.Fprintf(os.Stderr, "FATAL: db.host not provided! Exiting...\n")
        os.Exit(1)
    }
    if app.Config.DB.Name == "" {
        fmt.Fprintf(os.Stderr, "FATAL: db.name not provided! Exiting...\n")
        os.Exit(1)
    }
    if app.Config.DB.User == "" {
        fmt.Fprintf(os.Stderr, "FATAL: db.user not provided! Exiting...\n")
        os.Exit(1)
    }
    if app.Config.DB.Pass == "" {
        fmt.Fprintf(os.Stderr, "FATAL: db.pass not provided! Exiting...\n")
        os.Exit(1)
    }

    var err error
    app.DB, err = sqlx.Connect(
        "postgres",
        fmt.Sprintf(
            "host=%s port=%d user=%s dbname=%s password='%s' sslmode=disable",
            app.Config.DB.Host,
            app.Config.DB.Port,
            app.Config.DB.User,
            app.Config.DB.Name,
            app.Config.DB.Pass,
        ),
    )
    if err != nil {
        fmt.Fprintf(os.Stderr, "FATAL: Unable to initialise database: %v\n", err)
        os.Exit(1)
    }
    app.DB.SetConnMaxLifetime(time.Minute * 3)
    app.DB.SetMaxOpenConns(10)
    app.DB.SetMaxIdleConns(10)
    defer app.DB.Close()

    // handle command arguments
    if len(os.Args) > 1 {
        arg := os.Args[1]

        switch arg {
        case "createTOTP":
            if len(os.Args) < 4 {
                fmt.Fprintf(os.Stderr, "FATAL: `username` and `name` must be specified for createTOTP.\n")
                os.Exit(1)
            }
            username := os.Args[2]
            totpName := os.Args[3]

            account, err := controller.GetAccountByUsername(app.DB, username)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to fetch account \"%s\": %v\n", username, err)
                os.Exit(1)
            }

            if account == nil {
                fmt.Fprintf(os.Stderr, "FATAL: Account \"%s\" does not exist.\n", username)
                os.Exit(1)
            }

            secret := controller.GenerateTOTPSecret(controller.TOTP_SECRET_LENGTH)
            totp := model.TOTP {
                AccountID: account.ID,
                Name: totpName,
                Secret: string(secret),
            }

            err = controller.CreateTOTP(app.DB, &totp)
            if err != nil {
                if strings.HasPrefix(err.Error(), "pq: duplicate key") {
                    fmt.Fprintf(os.Stderr, "FATAL: Account \"%s\" already has a TOTP method named \"%s\"!\n", account.Username, totp.Name)
                    os.Exit(1)
                }
                fmt.Fprintf(os.Stderr, "FATAL: Failed to create TOTP method: %v\n", err)
                os.Exit(1)
            }

            url := controller.GenerateTOTPURI(account.Username, totp.Secret)
            fmt.Printf("%s\n", url)
            return

        case "deleteTOTP":
            if len(os.Args) < 4 {
                fmt.Fprintf(os.Stderr, "FATAL: `username` and `name` must be specified for deleteTOTP.\n")
                os.Exit(1)
            }
            username := os.Args[2]
            totpName := os.Args[3]

            account, err := controller.GetAccountByUsername(app.DB, username)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to fetch account \"%s\": %v\n", username, err)
                os.Exit(1)
            }

            if account == nil {
                fmt.Fprintf(os.Stderr, "FATAL: Account \"%s\" does not exist.\n", username)
                os.Exit(1)
            }

            err = controller.DeleteTOTP(app.DB, account.ID, totpName)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to create TOTP method: %v\n", err)
                os.Exit(1)
            }

            fmt.Printf("TOTP method \"%s\" deleted.\n", totpName)
            return

        case "listTOTP":
            if len(os.Args) < 3 {
                fmt.Fprintf(os.Stderr, "FATAL: `username` must be specified for listTOTP.\n")
                os.Exit(1)
            }
            username := os.Args[2]

            account, err := controller.GetAccountByUsername(app.DB, username)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to fetch account \"%s\": %v\n", username, err)
                os.Exit(1)
            }

            if account == nil {
                fmt.Fprintf(os.Stderr, "FATAL: Account \"%s\" does not exist.\n", username)
                os.Exit(1)
            }

            totps, err := controller.GetTOTPsForAccount(app.DB, account.ID)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to create TOTP methods: %v\n", err)
                os.Exit(1)
            }

            for i, totp := range totps {
                fmt.Printf("%d. %s - Created %s\n", i + 1, totp.Name, totp.CreatedAt)
            }
            if len(totps) == 0 {
                fmt.Printf("\"%s\" has no TOTP methods.\n", account.Username)
            }
            return

        case "testTOTP":
            if len(os.Args) < 4 {
                fmt.Fprintf(os.Stderr, "FATAL: `username` and `name` must be specified for testTOTP.\n")
                os.Exit(1)
            }
            username := os.Args[2]
            totpName := os.Args[3]

            account, err := controller.GetAccountByUsername(app.DB, username)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to fetch account \"%s\": %v\n", username, err)
                os.Exit(1)
            }

            if account == nil {
                fmt.Fprintf(os.Stderr, "FATAL: Account \"%s\" does not exist.\n", username)
                os.Exit(1)
            }

            totp, err := controller.GetTOTP(app.DB, account.ID, totpName)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to fetch TOTP method \"%s\": %v\n", totpName, err)
                os.Exit(1)
            }

            if totp == nil {
                fmt.Fprintf(os.Stderr, "FATAL: TOTP method \"%s\" does not exist for account \"%s\"\n", totpName, username)
                os.Exit(1)
            }

            code := controller.GenerateTOTP(totp.Secret, 0)
            fmt.Printf("%s\n", code)
            return
        
        case "cleanTOTP":
            err := controller.DeleteUnconfirmedTOTPs(app.DB)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to clean up TOTP methods: %v\n", err)
                os.Exit(1)
            }
            fmt.Printf("Cleaned up dangling TOTP methods successfully.\n")
            return

        case "createInvite":
            fmt.Printf("Creating invite...\n")
            invite, err := controller.CreateInvite(app.DB, 16, time.Hour * 24)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to create invite code: %v\n", err)
                os.Exit(1)
            }

            fmt.Printf(
                "Here you go! This code expires in %d hours: %s\n",
                int(math.Ceil(invite.ExpiresAt.Sub(invite.CreatedAt).Hours())),
                invite.Code,
            )
            return

        case "purgeInvites":
            fmt.Printf("Deleting all invites...\n")
            err := controller.DeleteAllInvites(app.DB)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to delete invites: %v\n", err)
                os.Exit(1)
            }

            fmt.Printf("Invites deleted successfully.\n")
            return

        case "listAccounts":
            accounts, err := controller.GetAllAccounts(app.DB)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to fetch accounts: %v\n", err)
                os.Exit(1)
            }

            for _, account := range accounts {
                email := "<none>"
                if account.Email.Valid { email = account.Email.String }
                fmt.Printf(
                    "User: %s\n" +
                    "\tID: %s\n" +
                    "\tEmail: %s\n" +
                    "\tCreated: %s\n",
                    account.Username,
                    account.ID,
                    email,
                    account.CreatedAt,
                )
            }
            return

        case "changePassword":
            if len(os.Args) < 4 {
                fmt.Fprintf(os.Stderr, "FATAL: `username` and `password` must be specified for changePassword\n")
                os.Exit(1)
            }

            username := os.Args[2]
            password := os.Args[3]
            account, err := controller.GetAccountByUsername(app.DB, username)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to fetch account \"%s\": %v\n", username, err)
                os.Exit(1)
            }
            if account == nil {
                fmt.Fprintf(os.Stderr, "FATAL: Account \"%s\" does not exist.\n", username)
                os.Exit(1)
            }

            hashedPassword, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to update password: %v\n", err)
                os.Exit(1)
            }
            account.Password = string(hashedPassword)
            err = controller.UpdateAccount(app.DB, account)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to delete account: %v\n", err)
                os.Exit(1)
            }

            fmt.Printf("Account \"%s\" deleted successfully.\n", account.Username)
            return

        case "deleteAccount":
            if len(os.Args) < 3 {
                fmt.Fprintf(os.Stderr, "FATAL: `username` must be specified for deleteAccount\n")
                os.Exit(1)
            }
            username := os.Args[2]
            fmt.Printf("Deleting account \"%s\"...\n", username)

            account, err := controller.GetAccountByUsername(app.DB, username)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to fetch account \"%s\": %v\n", username, err)
                os.Exit(1)
            }

            if account == nil {
                fmt.Fprintf(os.Stderr, "FATAL: Account \"%s\" does not exist.\n", username)
                os.Exit(1)
            }

            fmt.Printf("You are about to delete \"%s\". Are you sure? (y/[N]): ", account.Username)
            res := ""
            fmt.Scanln(&res)
            if !strings.HasPrefix(res, "y") {
                return
            }
            
            err = controller.DeleteAccount(app.DB, account.ID)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to delete account: %v\n", err)
                os.Exit(1)
            }

            fmt.Printf("Account \"%s\" deleted successfully.\n", account.Username)
            return

        case "testLogSearch":
            // TODO: rename to "logs"; add parameters
            logger := log.Logger { DB: app.DB }
            logs, err := logger.Search([]log.LogLevel{ log.LEVEL_INFO, log.LEVEL_WARN }, []string{ log.TYPE_ACCOUNT, log.TYPE_MUSIC }, "ari", 0, 100)
            if err != nil {
                fmt.Fprintf(os.Stderr, "FATAL: Failed to fetch logs: %v\n", err)
                os.Exit(1)
            }
            for _, log := range(logs) {
                fmt.Printf("[%s] [%s] [%d] [%s] %s\n", log.CreatedAt.Format(time.UnixDate), log.ID, log.Level, log.Type, log.Content)
            }
            return
        }

        // command help
        fmt.Print(
            "Available commands:\n\n" +
            "createTOTP <username> <name>:\n\tCreates a timed one-time passcode method.\n" +
            "listTOTP <username>:\n\tLists an account's TOTP methods.\n" +
            "deleteTOTP <username> <name>:\n\tDeletes an account's TOTP method.\n" +
            "testTOTP <username> <name>:\n\tGenerates the code for an account's TOTP method.\n" +
            "cleanTOTP:\n\tCleans up unconfirmed (dangling) TOTP methods.\n" +
            "\n" +
            "createInvite:\n\tCreates an invite code to register new accounts.\n" +
            "purgeInvites:\n\tDeletes all available invite codes.\n" +
            "listAccounts:\n\tLists all active accounts.\n",
            "deleteAccount <username>:\n\tDeletes an account with a given `username`.\n",
        )
        return
    }

    // handle DB migrations
    controller.CheckDBVersionAndMigrate(app.DB)

    // initial invite code
    accountsCount := 0
    err = app.DB.Get(&accountsCount, "SELECT count(*) FROM account")
    if err != nil { panic(err) }
    if accountsCount == 0 {
        _, err := app.DB.Exec("DELETE FROM invite")
        if err != nil {
            fmt.Fprintf(os.Stderr, "FATAL: Failed to clear existing invite codes: %v\n", err)
            os.Exit(1)
        }

        invite, err := controller.CreateInvite(app.DB, 16, time.Hour * 24)
        if err != nil {
            fmt.Fprintf(os.Stderr, "FATAL: Failed to create invite code: %v\n", err)
            os.Exit(1)
        }

        fmt.Printf("No accounts exist! Generated invite code: %s\n", invite.Code)
    }

    // delete expired invites
    err = controller.DeleteExpiredInvites(app.DB)
    if err != nil {
        fmt.Fprintf(os.Stderr, "FATAL: Failed to clear expired invite codes: %v\n", err)
        os.Exit(1)
    }

    // clean up unconfirmed TOTP methods
    err = controller.DeleteUnconfirmedTOTPs(app.DB)
    if err != nil {
        fmt.Fprintf(os.Stderr, "FATAL: Failed to clean up unconfirmed TOTP methods: %v\n", err)
        os.Exit(1)
    }

    // start the web server!
    mux := createServeMux(&app)
    fmt.Printf("Now serving at http://%s:%d\n", app.Config.Host, app.Config.Port)
    stdLog.Fatal(
        http.ListenAndServe(fmt.Sprintf("%s:%d", app.Config.Host, app.Config.Port), 
        HTTPLog(DefaultHeaders(mux)),
    ))
}

func createServeMux(app *model.AppState) *http.ServeMux {
    mux := http.NewServeMux()

    mux.Handle("/admin/", http.StripPrefix("/admin", admin.Handler(app)))
    mux.Handle("/api/", http.StripPrefix("/api", api.Handler(app)))
    mux.Handle("/music/", http.StripPrefix("/music", view.MusicHandler(app)))
    mux.Handle("/uploads/", http.StripPrefix("/uploads", staticHandler(filepath.Join(app.Config.DataDirectory, "uploads"))))
    mux.Handle("/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        if r.Method == http.MethodHead {
            w.WriteHeader(http.StatusOK)
            return
        }

        if r.URL.Path == "/" || r.URL.Path == "/index.html" {
            err := templates.IndexTemplate.Execute(w, nil)
            if err != nil {
                http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
            }
            return
        }
        staticHandler("public").ServeHTTP(w, r)
    }))

    return mux
}

func staticHandler(directory string) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        info, err := os.Stat(filepath.Join(directory, filepath.Clean(r.URL.Path)))

        // does the file exist?
        if err != nil {
            if errors.Is(err, os.ErrNotExist) {
                http.NotFound(w, r)
                return
            }
        }

        // is thjs a directory? (forbidden)
        if info.IsDir() {
            http.NotFound(w, r)
            return
        }

        http.FileServer(http.Dir(directory)).ServeHTTP(w, r)
    })
}

var PoweredByStrings = []string{
    "nerd rage",
    "estrogen",
    "your mother",
    "awesome powers beyond comprehension",
    "jared",
    "the weight of my sins",
    "the arc reactor",
    "AA batteries",
    "15 euro solar panel from ebay",
    "magnets, how do they work",
    "a fax machine",
    "dell optiplex",
    "a trans girl's nintendo wii",
    "BASS",
    "electricity, duh",
    "seven hamsters in a big wheel",
    "girls",
    "mzungu hosting",
    "golang",
    "the state of the world right now",
    "Mozilla/4.0 (compatible; MSIE 6.0; Windows NT 5.1; SV1)",
    "the good folks at aperture science",
    "free2play CDs",
    "aridoodle",
    "the love of creating",
    "not for the sake of art; not for the sake of money; we like painting naked people",
    "30 billion dollars in VC funding",
}

func DefaultHeaders(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        w.Header().Add("Server", "arimelody.me")
        w.Header().Add("Do-Not-Stab", "1")
        w.Header().Add("X-Clacks-Overhead", "GNU Terry Pratchett")
        w.Header().Add("X-Hacker", "spare me please")
        w.Header().Add("X-Robots-TXT", "'; DROP TABLE pages;")
        w.Header().Add("X-Thinking-With", "Portals")
        w.Header().Add(
            "X-Powered-By",
            PoweredByStrings[rand.Intn(len(PoweredByStrings))],
        )
        next.ServeHTTP(w, r)
    })
}

type LoggingResponseWriter struct {
    http.ResponseWriter
    Status int
}

func (lrw *LoggingResponseWriter) WriteHeader(status int) {
    lrw.Status = status
    lrw.ResponseWriter.WriteHeader(status)
}

func HTTPLog(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        start := time.Now()

        lrw := LoggingResponseWriter{w, http.StatusOK}

        next.ServeHTTP(&lrw, r)

        after := time.Now()
        difference := (after.Nanosecond() - start.Nanosecond()) / 1_000_000
        elapsed := "<1"
        if difference >= 1 {
            elapsed = strconv.Itoa(difference)
        }

        statusColour := colour.Reset

        if lrw.Status - 600 <= 0 { statusColour = colour.Red }
        if lrw.Status - 500 <= 0 { statusColour = colour.Yellow }
        if lrw.Status - 400 <= 0 { statusColour = colour.White }
        if lrw.Status - 300 <= 0 { statusColour = colour.Green }

        fmt.Printf("[%s] %s %s - %s%d%s (%sms) (%s)\n",
        after.Format(time.UnixDate),
        r.Method,
        r.URL.Path,
        statusColour,
        lrw.Status,
        colour.Reset,
        elapsed,
        r.Header["User-Agent"][0])
    })
}