slink/slink.go

package main

import (
  "log"
  "net/http"
  "bufio"
  "io/ioutil"
  "encoding/json"
  "time"
  "os"
  "math/rand"
  "unsafe"
  "strings"
  "flag"
)

const (
  letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
  letterIdxBits = 6
  letterIdxMask = 1<<letterIdxBits - 1
  letterIdxMax  = 63 / letterIdxBits
)

var src = rand.NewSource(time.Now().UnixNano())

type Url struct{
  Long string `json:"url"`
  Timestamp int64 `json:"timestamp"`
  User string `json:"user"`
}
type Link struct{
  Long string `json:"url"`
  Timestamp int64 `json:"timestamp"`
}

var StorageMap = make(map[string]Url)

func RandString(n int) string {
  b := make([]byte, n)
  for i, cache, remain := n-1, src.Int63(), letterIdxMax; i >= 0; {
    if remain == 0 {
      cache, remain = src.Int63(), letterIdxMax
    }
    if idx := int(cache & letterIdxMask); idx < len(letterBytes) {
      b[i] = letterBytes[idx]
      i--
    }
    cache >>= letterIdxBits
    remain--
  }
  return *(*string)(unsafe.Pointer(&b))
}

func logSetup(port string,base string,storage string,usersFile string,host string) {
  log.Printf("Server will run on: %s\n",port)
  log.Printf("Server will use: %s\n", base)
  log.Printf("Server will use: %s\n", storage)
  log.Printf("Reading users from: %s\n",usersFile)
  log.Printf("Using hostname: %s\n",host)
}

func healthHandler(w http.ResponseWriter, r *http.Request){
  w.WriteHeader(200)
}

func getToken(token string, usersFile string) (user string) {
  user = ""
  file, err := os.Open(usersFile)
  if err != nil {
    return ""
  }
  defer file.Close()
  scanner := bufio.NewScanner(file)
  for scanner.Scan() {
    splitstring := strings.Split(scanner.Text(),"|")
    if splitstring[1] == token {
      user = splitstring[0]
    }
  }
  return user
}

func main() {
  Port := flag.String("port","8080","Port to listen on")
  base := flag.String("base","/","Basepath /example/")
  storage := flag.String("storage","storage.json","Storage json file")
  usersFile := flag.String("users","users.txt","file for users and keys")
  proto := flag.String("proto","https","https | http")
  host := flag.String("host","localhost","hostname")
  flag.Parse()
  logSetup(*Port, *base, *storage, *usersFile, *host)
  if _, err := os.Stat(*storage); err == nil {
    log.Println("Loading data from storage file")
    data,_ := ioutil.ReadFile (*storage)
    json.Unmarshal(data,&StorageMap)
  }
  log.Printf("%+v",StorageMap)
  ListenAddress := ":"+*Port
  http.HandleFunc(*base, func (w http.ResponseWriter, r *http.Request){
    short := strings.TrimPrefix(r.RequestURI,*base)
    log.Printf(short)
    var long string
    if val,ok := StorageMap[short]; ok {
      long = val.Long
    }else{
      w.WriteHeader(404)
      return
    }
    http.Redirect(w,r,long,301)
  })
  http.HandleFunc(*base+"health",healthHandler)
  http.HandleFunc(*base+"admin/create", func (w http.ResponseWriter, r *http.Request){
    authHeader,hasAuthHeader := r.Header["Authentication"]
    if hasAuthHeader {
      user := getToken(authHeader[0],*usersFile)
      if user != ""{
        log.Printf("user authed: %s",user)
        body,_ := ioutil.ReadAll(r.Body)
        defer r.Body.Close()
        var rq Url
        err := json.Unmarshal(body,&rq)
        if err != nil {
          log.Printf(err.Error())
          w.WriteHeader(500)
          w.Write([]byte("Json parse error"))
          return
        }
        if rq.Long == "" {
          w.WriteHeader(500)
          w.Write([]byte("No Url provided"))
          return
        }
        rq.Timestamp = time.Now().Unix()
        rq.User = user
        short := RandString(8)
        log.Println(short)
        StorageMap[short] = rq
        var rrr strings.Builder
        rrr.WriteString(*proto)
        rrr.WriteString("://")
        rrr.WriteString(*host)
        rrr.WriteString(*base)
        rrr.WriteString(short)
        rrr.WriteString("\n")
        w.Write([]byte(rrr.String()))
        jsondump,_ := json.Marshal(StorageMap)
        ioutil.WriteFile(*storage,jsondump,0600)
      }else{
        w.WriteHeader(401)
      }
    }else{
      w.WriteHeader(401)
    }
  })
  http.HandleFunc(*base+"admin/list",func (w http.ResponseWriter, r *http.Request){
    authHeader,hasAuthHeader := r.Header["Authentication"]
    if hasAuthHeader {
      user := getToken(authHeader[0],*usersFile)
      log.Printf("user authed: %s",user)
      if user != ""{
        userMap := make(map[string]Link)
        for key,link := range StorageMap {
         if link.User == user {
           userMap[key] = Link{
             Long: link.Long,
             Timestamp: link.Timestamp,
           }
         }
        }
        jsondump,_ := json.Marshal(userMap)
        w.Write([]byte(jsondump))
      }else{
        w.WriteHeader(401)
      }
    }else{
      w.WriteHeader(401)
    }
  })
  http.HandleFunc(*base+"admin/del",func (w http.ResponseWriter, r *http.Request){
    authHeader,hasAuthHeader := r.Header["Authentication"]
    if hasAuthHeader {
      user := getToken(authHeader[0],*usersFile)
      if user != ""{
        log.Printf("user authed: %s",user)
        body,_ := ioutil.ReadAll(r.Body)
        defer r.Body.Close()
        var rq Url
        json.Unmarshal(body,&rq)
        log.Printf(rq.Long)
        if StorageMap[rq.Long].User == user {
          delete(StorageMap,rq.Long)
        }else{
          w.WriteHeader(404)
          return
        }
        w.WriteHeader(200)
        jsondump,_ := json.Marshal(StorageMap)
        ioutil.WriteFile(*storage,jsondump,0600)
      }else{
        w.WriteHeader(401)
      }
    }else{
      w.WriteHeader(401)
    }
  })
  if err := http.ListenAndServe(ListenAddress, nil); err != nil {
    panic(err)
  }
}