sus/main.go

// sus - a simple url shortener Copyright (C) 2022 Akbar Rahman (hi@alv.cx)
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

package main

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/go-redis/redis/v8"
	"github.com/gorilla/mux"
	"golang.org/x/net/context"
)

var client = redis.NewClient(&redis.Options{
	Addr:     "redis:6379",
	Password: "",
	DB:       0,
})

var SECRET string
var INDEX_GET_REDIRECT = "http://alv.cx"

func main() {
	r := mux.NewRouter()
	r.HandleFunc("/{shortlink}", shortlinkHandler)
	r.HandleFunc("/", indexHandler)

	listenAddress := "0.0.0.0:80"

	if p, ok := os.LookupEnv("LISTEN_ADDRESS"); ok {
		listenAddress = p
	}

	if p, ok := os.LookupEnv("SECRET"); ok {
		SECRET = p
	}

	if p, ok := os.LookupEnv("INDEX_GET_REDIRECT"); ok {
		INDEX_GET_REDIRECT = p
	}

	log.Fatal(http.ListenAndServe(listenAddress, r))
}

func indexHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Println("indexHandler called")
	if r.Method != "POST" {
		http.Redirect(w, r, INDEX_GET_REDIRECT, 302)
		return
	}

	r.ParseForm()

	command := r.PostForm.Get("Command")
	shortlink := r.PostForm.Get("Shortlink")
	redirect := r.PostForm.Get("Redirect")
	alt_redirect := r.PostForm.Get("AltRedirect")
	alt_condition := r.PostForm.Get("AltCondition")
	fmt.Printf("command: %v, shortlink: %v, redirect: %v, alt_redirect: %v, alt_condition: %v\n",
		command,
		shortlink,
		redirect,
		alt_redirect,
		alt_condition,
	)

	formstring := command + ":" + shortlink + ":" + redirect + ":" + alt_condition + ":" + alt_redirect

	fmt.Println("formstring: " + formstring)

	signature := r.Header.Get("Signature")
	calculatedSignature := fmt.Sprintf(
		"SUS-SIGNATURE-%v",
		getSha256HMACSignature(
			[]byte(SECRET),
			command+":"+shortlink+":"+redirect+":"+alt_condition+":"+alt_redirect,
		),
	)

	if signature != calculatedSignature {
		fmt.Println("signature do no match")
		fmt.Println(signature)
		fmt.Println(calculatedSignature)
		w.WriteHeader(401)
		w.Write([]byte("401 Unauthorized"))
		return
	}

	if command == "create" {

		ctx := context.Background()
		_, err := client.Get(ctx, shortlink).Result()
		if err == redis.Nil {
			err = client.Set(ctx, shortlink+":redirect", shortlink, 0).Err()
			if err != nil {
				fmt.Println(err)
				w.WriteHeader(500)
				w.Write([]byte("500 Internal Server Error"))
				return
			}
			err = client.Set(ctx, shortlink+":altcondition", alt_condition, 0).Err()
			if err != nil {
				fmt.Println(err)
				w.WriteHeader(500)
				w.Write([]byte("500 Internal Server Error"))
				return
			}
			err = client.Set(ctx, shortlink+":altredirect", alt_redirect, 0).Err()

			if err != nil {
				fmt.Println(err)
				w.WriteHeader(500)
				w.Write([]byte("500 Internal Server Error"))
				return
			}

			w.WriteHeader(200)
			w.Write([]byte("200 Success"))

			return

		} else if err != nil {
			fmt.Println(err)
			w.WriteHeader(500)
			w.Write([]byte("500 Internal Server Error"))
			return
		}

		fmt.Println(err)
		w.WriteHeader(403)
		w.Write([]byte("403 Forbidden"))
		return
	}

	if command == "delete" {
		ctx := context.Background()
		if err := client.Del(ctx, shortlink+":redirect").Err(); err != nil {
			w.WriteHeader(500)
			w.Write([]byte("500 Internal Server Error"))
		}
		if err := client.Del(ctx, shortlink+":altredirect").Err(); err != nil {
			w.WriteHeader(500)
			w.Write([]byte("500 Internal Server Error"))
		}
		if err := client.Del(ctx, shortlink+":altcondition").Err(); err != nil {
			w.WriteHeader(500)
			w.Write([]byte("500 Internal Server Error"))
		}
	}

	if command == "list" {
		ctx := context.Background()
		keys, err := client.Keys(ctx, "*").Result()
		if err != nil {
			fmt.Println(err)
			w.WriteHeader(500)
			w.Write([]byte("500 Internal Server Error"))
			return
		}

		resp := ""
		for _, key := range keys {
			shortlink, err = client.Get(ctx, key).Result()
			if err == redis.Nil {
				w.WriteHeader(500)
				w.Write([]byte("500 Internal Server Error"))
				return
			}
			resp += key + ":" + shortlink + "\n"
		}

		w.WriteHeader(200)
		w.Write([]byte(resp))
	}

}

func shortlinkHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Println("shortlinkHandler called")
	shortlink := string(mux.Vars(r)["shortlink"])
	ctx := context.Background()
	client.Incr(ctx, shortlink + ":hits")

	redirect, err := client.Get(ctx, shortlink+":redirect").Result()
	fmt.Printf("shortlink: %v, redirect: %v\n", shortlink, redirect)
	if err == redis.Nil {
		w.WriteHeader(404)
		w.Write([]byte("404 Not Found"))
		return
	} else if err != nil {
		fmt.Println(err)
		fmt.Println(0)
		w.WriteHeader(500)
		w.Write([]byte("500 Internal Server Error"))
		return
	}

	altcondition, err := client.Get(ctx, shortlink+":altcondition").Result()
	if err == redis.Nil {
		http.Redirect(w, r, redirect, 302)
		return
	} else if err != nil {
		fmt.Println(err)
		fmt.Println(1)
		w.WriteHeader(500)
		w.Write([]byte("500 Internal Server Error"))
		return
	}

	altredirect, err := client.Get(ctx, shortlink+":altredirect").Result()
	if err == redis.Nil {
		http.Redirect(w, r, redirect, 302)
		return
	} else if err != nil {
		fmt.Println(err)
		fmt.Println(2)
		w.WriteHeader(500)
		w.Write([]byte("500 Internal Server Error"))
		return
	}

	altcondition_split := strings.Split(altcondition, ",")
	ac_varname := altcondition_split[0]
	ac_operator := altcondition_split[1]
	ac_required_value, _ := strconv.Atoi(altcondition_split[2])

	var ac_varval int

	if ac_varname != "timestamp" {
		ac_varvalstr, err := client.Get(ctx, shortlink+":"+ac_varname).Result()
		if err == redis.Nil {
			ac_varval = 0
		} else if err != nil {
			fmt.Println(err)
			fmt.Println(3)
			w.WriteHeader(500)
			w.Write([]byte("500 Internal Server Error"))
			return
		} else {
			ac_varval, _ = strconv.Atoi(ac_varvalstr)
		}
	} else {
		ac_varval = int(time.Now().Unix())
	}

	if (ac_operator == "eq" && ac_varval == ac_required_value) ||
		(ac_operator == "gt" && ac_varval > ac_required_value) ||
		(ac_operator == "lt" && ac_varval < ac_required_value) {
		http.Redirect(w, r, altredirect, 307)
	} else {
		http.Redirect(w, r, redirect, 307)
	}
}

func getSha256HMACSignature(secret []byte, data string) string {
	h := hmac.New(sha256.New, secret)
	io.WriteString(h, data)
	return hex.EncodeToString(h.Sum(nil))
}