package main

import (
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path"
	"strings"
	"sync"
	"syscall"
	"time"
)

var (
	S3_HOST          string
	S3_SECRET_KEY    = EnvString("S3_SECRET_KEY", "xyz")
	S3_ACCESS_KEY    = EnvString("S3_ACCESS_KEY", "123")
	S3_ENDPOINT      = EnvString("S3_ENDPOINT", "https://bucket.s3.region.host.tld")
	S3_REGION        = EnvString("S3_REGION", "region")
	S3_BUCKET        = EnvString("S3_BUCKET", "bucket")
	HTTP_KEY         = EnvString("HTTP_KEY", "teto")
	HTTP_ADDRESS     = EnvString("HTTP_ADDRESS", "localhost:8080")
	HTTP_CORS_ORIGIN = EnvString("HTTP_CORS_ORIGIN", "")
)

// ----------------------------------------------------------------------------
// HELPER FUNCTIONS
// ----------------------------------------------------------------------------

func EnvString(field, initial string) string {
	if value := os.Getenv(field); value == "" {
		return initial
	} else {
		return value
	}
}

func SHA256HEX(data []byte) string {
	h := sha256.Sum256(data)
	return hex.EncodeToString(h[:])
}

func SHA256HMAC(key, data []byte) []byte {
	h := hmac.New(sha256.New, key)
	h.Write(data)
	return h.Sum(nil)
}

func ValidateToken(r *http.Request) bool {
	requestPath := r.URL.Path
	requestToken := r.URL.Query().Get("token")
	if requestToken == "" {
		return false
	}

	// Decode Token Segments
	s := strings.SplitN(requestToken, ".", 2)
	if len(s) != 2 {
		return false
	}
	payload, err := base64.RawURLEncoding.DecodeString(s[0])
	if err != nil {
		return false
	}
	signature, err := base64.RawURLEncoding.DecodeString(s[1])
	if err != nil {
		return false
	}

	// Validate Token Signature
	hash := hmac.New(sha256.New, []byte(HTTP_KEY))
	hash.Write(payload)
	if !hmac.Equal(signature, hash.Sum(nil)) {
		return false
	}

	// Parse Token Data
	var TokenData struct {
		Pattern string `json:"pattern"`
		Expires int64  `json:"expires"`
	}
	if err := json.Unmarshal(payload, &TokenData); err != nil {
		return false
	}
	if time.Now().Unix() > TokenData.Expires {
		return false
	}
	if ok, err := path.Match(TokenData.Pattern, requestPath); err != nil || !ok {
		return false
	}

	// LGTM
	return true
}

func SignRequest(r *http.Request) {
	// Timestamp
	t := time.Now().UTC()
	dateStamp := t.Format("20060102")
	dataAmazon := t.Format("20060102T150405Z")

	// Hash Content
	payload := []byte{}
	payloadHashHexSHA := SHA256HEX(payload)

	// Apply Headers
	r.Header.Set("Host", S3_HOST)
	r.Header.Set("x-amz-date", dataAmazon)
	r.Header.Set("x-amz-content-sha256", payloadHashHexSHA)
	r.Host = S3_HOST

	// 1. Create a canonical request
	signedHeaders := "host;x-amz-content-sha256;x-amz-date"
	canonicalHeaders := fmt.Sprintf(
		"host:%s\nx-amz-content-sha256:%s\nx-amz-date:%s\n",
		r.Host, payloadHashHexSHA, dataAmazon,
	)
	canonicalRequest := fmt.Sprintf("%s\n%s\n%s\n%s\n%s\n%s",
		r.Method,
		r.URL.EscapedPath(),
		r.URL.Query().Encode(),
		canonicalHeaders,
		signedHeaders,
		payloadHashHexSHA,
	)

	// Step 2: String to sign
	service := "s3"
	scope := fmt.Sprintf("%s/%s/%s/aws4_request", dateStamp, S3_REGION, service)
	stringToSign := fmt.Sprintf("AWS4-HMAC-SHA256\n%s\n%s\n%s",
		dataAmazon, scope, SHA256HEX([]byte(canonicalRequest)),
	)

	// Step 3: Derive signing key
	kDate := SHA256HMAC([]byte("AWS4"+S3_SECRET_KEY), []byte(dateStamp))
	kRegion := SHA256HMAC(kDate, []byte(S3_REGION))
	kService := SHA256HMAC(kRegion, []byte(service))
	kSigning := SHA256HMAC(kService, []byte("aws4_request"))

	// Step 4: Signature
	signature := hex.EncodeToString(SHA256HMAC(kSigning, []byte(stringToSign)))
	authHeader := fmt.Sprintf(
		"AWS4-HMAC-SHA256 Credential=%s/%s, SignedHeaders=%s, Signature=%s",
		S3_ACCESS_KEY, scope, signedHeaders, signature,
	)
	r.Header.Set("Authorization", authHeader)
}

// ----------------------------------------------------------------------------
// SERVICE FUNCTIONS
// ----------------------------------------------------------------------------

func SetupMux() *http.ServeMux {
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {

		// Handle CORS Requests
		if HTTP_CORS_ORIGIN != "" {
			w.Header().Add("Access-Control-Allow-Origin", HTTP_CORS_ORIGIN)
			w.Header().Add("Access-Control-Allow-Methods", "GET, OPTIONS")
			w.Header().Add("Access-Control-Allow-Headers", "Content-Type")
			if r.Method == "OPTIONS" {
				w.WriteHeader(http.StatusNoContent)
				return
			}
		}

		// Request Validation
		if r.Method != "GET" {
			w.WriteHeader(http.StatusMethodNotAllowed)
			return
		}
		if r.URL.Path == "/" {
			w.WriteHeader(http.StatusNotFound)
			return
		}

		// Check Access Token for Private Endpoints
		if strings.HasPrefix(r.URL.Path, "/private") {
			w.Header().Add("Cache-Control", "no-store")
			if !ValidateToken(r) {
				w.WriteHeader(http.StatusUnauthorized)
				return
			}
		}

		// Request Bucket File
		url := fmt.Sprint(S3_ENDPOINT, r.URL.Path)
		req, err := http.NewRequest(http.MethodGet, url, http.NoBody)
		if err != nil {
			return
		}
		SignRequest(req)

		res, err := http.DefaultClient.Do(req)
		if err != nil {
			http.Error(w, "Upstream Error", http.StatusBadGateway)
			return
		}
		defer res.Body.Close()

		// Stream File Contents
		w.Header().Set("Cache-Control", "public, max-age=31536000, immutable")
		w.Header().Set("Content-Type", res.Header.Get("Content-Type"))
		w.WriteHeader(res.StatusCode)
		io.Copy(w, res.Body)
	})
	return mux
}

func StartupHTTP(stop context.Context, await *sync.WaitGroup) {
	var listener net.Listener
	var err error

	if strings.HasPrefix(HTTP_ADDRESS, "unix/") {
		path := strings.TrimPrefix(HTTP_ADDRESS, "unix/")
		os.Remove(path)
		listener, err = net.Listen("unix", path)
		if err == nil {
			os.Chmod(path, 0660)
		}
	} else {
		listener, err = net.Listen("tcp", HTTP_ADDRESS)
	}
	if err != nil {
		log.Fatalf("Listen Failed: %s\n", err)
		return
	}

	svr := http.Server{
		Handler:           SetupMux(),
		MaxHeaderBytes:    4096,
		IdleTimeout:       10 * time.Second,
		ReadHeaderTimeout: 10 * time.Second,
		ReadTimeout:       30 * time.Second,
		WriteTimeout:      30 * time.Second,
	}

	// Shutdown Logic
	await.Add(1)
	go func() {
		defer await.Done()
		<-stop.Done()

		shutdownCtx, cancel := context.WithTimeout(context.Background(), time.Minute)
		defer cancel()

		if err := svr.Shutdown(shutdownCtx); err != nil {
			log.Printf("Shutdown Error: %s\n", err)
			return
		}
		log.Println("Server Closed")
	}()

	log.Printf("Listening @ %s\n", HTTP_ADDRESS)
	if err := svr.Serve(listener); err != http.ErrServerClosed {
		log.Printf("Startup Failed: %s\n", err)
		return
	}
}

func main() {
	time.Local = time.UTC

	// Startup Services
	var stopCtx, stop = context.WithCancel(context.Background())
	var stopWg sync.WaitGroup
	go StartupHTTP(stopCtx, &stopWg)

	// Await Shutdown Signal
	cancel := make(chan os.Signal, 1)
	signal.Notify(cancel, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
	<-cancel
	stop()

	// Begin Shutdown Process
	log.Println("Shutting Down!")
	timeout, finish := context.WithTimeout(context.Background(), time.Minute)
	defer finish()
	go func() {
		<-timeout.Done()
		if timeout.Err() == context.DeadlineExceeded {
			log.Fatalln("Shutdown Deadline Exceeded")
			return
		}
	}()
	stopWg.Wait()
	os.Exit(0)
}