<commit_before>package main

import (
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"koding/virt"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"koding/migrators/useroverlay/token"

	"github.com/gorilla/mux"
	"github.com/koding/logging"
	"labix.org/v2/mgo/bson"
)

var (
	log      = logging.NewLogger("useroverlay")
	mode     = flag.String("mode", "serve", "mode can be serve or token")
	certFile = flag.String("cert", "certs/user_file_exporter_self_signed_cert.pem", "TLS cert file")
	keyFile  = flag.String("key", "certs/user_file_exporter_self_signed_key.pem", "TLS key file")
	vm       = flag.String("vm", "", "the vm id")
	username = flag.String("username", "", "the username")
)

func main() {
	flag.Parse()
	switch *mode {
	case "serve":
		serve()
	case "token":
		createToken()
	}
}

func serve() {
	r := mux.NewRouter()
	r.HandleFunc("/export-files", exportFiles).Methods("POST")
	http.Handle("/", basicAuth(r))
	if err := http.ListenAndServeTLS(":3000", *certFile, *keyFile, nil); err != nil {
		log.Error(err.Error())
	}
}

func basicAuth(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if ok, err := checkAuth(w, r); !ok {
			if err != nil {
				log.Error(err.Error())
			}
			respond(w, 403, "Access denied")
			return
		}
		next.ServeHTTP(w, r)
	})
}

func checkAuth(w http.ResponseWriter, r *http.Request) (bool, error) {
	auth := r.Header.Get("Authorization")
	s := strings.SplitN(auth, " ", 2)
	if len(s) != 2 {
		return false, fmt.Errorf("Invalid authorization header: '%v'", auth)
	}

	b, err := base64.StdEncoding.DecodeString(s[1])
	if err != nil {
		return false, err
	}

	pair := strings.SplitN(string(b), ":", 2)
	if len(pair) != 2 {
		return false, fmt.Errorf("Invalid authorization header: '%v'", auth)
	}

	// validate the password:
	if pair[1] == token.StringToken(pair[0], r.PostFormValue("vm")) {
		return true, nil
	}
	return false, nil
}

func exportFiles(w http.ResponseWriter, r *http.Request) {
	vm, status, message, err := validateRequest(w, r)
	if err != nil {
		log.Error(err.Error())
	}

	if status != 200 {
		respond(w, status, message)
		return
	}

	archive, err := exportUserFiles(vm)
	if err != nil {
		respond(w, 500, err.Error())
		return
	}
	w.Header().Set("Content-type", "application/octet-stream")
	if _, err := io.Copy(w, archive); err != nil {
		log.Error(err.Error())
	}
}

func validateRequest(w http.ResponseWriter, r *http.Request) (*virt.VM, int, string, error) {
	id := r.PostFormValue("vm")
	if !bson.IsObjectIdHex(id) {
		return nil, 400, "Bad input", nil
	}
	vm := &virt.VM{Id: bson.ObjectIdHex(id)}

	exists, err := vm.ExistsRBD()
	if err != nil {
		return vm, 500, "Internal error", err
	}
	if !exists {
		return vm, 404, "Not found", nil
	}
	if err := vm.LockRBD(); err != nil {
		return vm, 412, "RBD cannot be locked", err
	}
	defer func() {
		if err := vm.UnlockRBD(); err != nil {
			log.Error(err.Error())
		}
	}()
	return vm, 200, "", nil
}

func exportUserFiles(vm *virt.VM) (io.Reader, error) {
	// map the RBD:
	if out, err := exec.Command("/usr/bin/rbd", "map", "--pool", virt.VMPool, "--image", vm.String()).CombinedOutput(); err != nil {
		return nil, commandError("rbd map failed.", err, out)
	}
	defer func() {
		if out, err := exec.Command("/usr/bin/rbd", "unmap", "/dev/rbd/"+virt.VMPool+"/"+vm.String()).CombinedOutput(); err != nil {
			log.Error(commandError("rbd unmap failed.", err, out).Error())
		}
	}()
	// generate a timestamp for uniqueness:<commit_msg>useroverlay: move rbd locking to exportFiles function<commit_after>package main

import (
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"koding/virt"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"koding/migrators/useroverlay/token"

	"github.com/gorilla/mux"
	"github.com/koding/logging"
	"labix.org/v2/mgo/bson"
)

var (
	log      = logging.NewLogger("useroverlay")
	mode     = flag.String("mode", "serve", "mode can be serve or token")
	certFile = flag.String("cert", "certs/user_file_exporter_self_signed_cert.pem", "TLS cert file")
	keyFile  = flag.String("key", "certs/user_file_exporter_self_signed_key.pem", "TLS key file")
	vm       = flag.String("vm", "", "the vm id")
	username = flag.String("username", "", "the username")
)

func main() {
	flag.Parse()
	switch *mode {
	case "serve":
		serve()
	case "token":
		createToken()
	}
}

func serve() {
	r := mux.NewRouter()
	r.HandleFunc("/export-files", exportFiles).Methods("POST")
	http.Handle("/", basicAuth(r))
	if err := http.ListenAndServeTLS(":3000", *certFile, *keyFile, nil); err != nil {
		log.Error(err.Error())
	}
}

func basicAuth(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if ok, err := checkAuth(w, r); !ok {
			if err != nil {
				log.Error(err.Error())
			}
			respond(w, 403, "Access denied")
			return
		}
		next.ServeHTTP(w, r)
	})
}

func checkAuth(w http.ResponseWriter, r *http.Request) (bool, error) {
	auth := r.Header.Get("Authorization")
	s := strings.SplitN(auth, " ", 2)
	if len(s) != 2 {
		return false, fmt.Errorf("Invalid authorization header: '%v'", auth)
	}

	b, err := base64.StdEncoding.DecodeString(s[1])
	if err != nil {
		return false, err
	}

	pair := strings.SplitN(string(b), ":", 2)
	if len(pair) != 2 {
		return false, fmt.Errorf("Invalid authorization header: '%v'", auth)
	}

	// validate the password:
	if pair[1] == token.StringToken(pair[0], r.PostFormValue("vm")) {
		return true, nil
	}
	return false, nil
}

func exportFiles(w http.ResponseWriter, r *http.Request) {
	vm, status, message, err := validateRequest(w, r)
	if err != nil {
		log.Error(err.Error())
	}

	if status != 200 {
		respond(w, status, message)
		return
	}
	if err := vm.LockRBD(); err != nil {
		respond(w, 412, "RBD cannot be locked")
	}
	defer func() {
		if err := vm.UnlockRBD(); err != nil {
			log.Error(err.Error())
		}
	}()
	archive, err := exportUserFiles(vm)
	if err != nil {
		respond(w, 500, err.Error())
		return
	}
	w.Header().Set("Content-type", "application/octet-stream")
	if _, err := io.Copy(w, archive); err != nil {
		log.Error(err.Error())
	}
}

func validateRequest(w http.ResponseWriter, r *http.Request) (*virt.VM, int, string, error) {
	id := r.PostFormValue("vm")
	if !bson.IsObjectIdHex(id) {
		return nil, 400, "Bad input", nil
	}
	vm := &virt.VM{Id: bson.ObjectIdHex(id)}

	exists, err := vm.ExistsRBD()
	if err != nil {
		return vm, 500, "Internal error", err
	}
	if !exists {
		return vm, 404, "Not found", nil
	}
	return vm, 200, "", nil
}

func exportUserFiles(vm *virt.VM) (io.Reader, error) {
	// map the RBD:
	if out, err := exec.Command("/usr/bin/rbd", "map", "--pool", virt.VMPool, "--image", vm.String()).CombinedOutput(); err != nil {
		return nil, commandError("rbd map failed.", err, out)
	}
	defer func() {
		if out, err := exec.Command("/usr/bin/rbd", "unmap", "/dev/rbd/"+virt.VMPool+"/"+vm.String()).CombinedOutput(); err != nil {
			log.Error(commandError("rbd unmap failed.", err, out).Error())
		}
	}()
	// generate a timestamp for uniqueness:
{'fraction_non_alphanumeric': 0.18841195636775748, 'fraction_numerical': 0.011333050007083156, 'mean_word_length': 3.8622589531680442, 'pattern_counts': {'":': 8, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2884823', 'n_tokens_mistral': 2833, 'n_tokens_neox': 2672, 'n_words': 637}