<commit_before>package main

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/kr/binarydist"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"
)

const (
	upcktimePath = "cktime"
	plat         = runtime.GOOS + "-" + runtime.GOARCH
)

const devValidTime = 7 * 24 * time.Hour

// Update protocol.
//
//   GET hk.heroku.com/hk-current-linux-amd64.json
//
//   200 ok
//   {
//       "Version": "2",
//       "Sha256": "..." // base64
//   }
//
// then
//
//   GET hkpatch.s3.amazonaws.com/hk-1-linux-amd64-to-2
//
//   200 ok
//   [bsdiff data]
//
// or
//
//   GET hkdist.s3.amazonaws.com/hk-2-linux-amd64.gz
//
//   200 ok
//   [gzipped executable data]
type Updater struct {
	hkURL   string
	binURL  string
	diffURL string
	dir     string
	info    struct {
		Version   string
		Sha256 []byte
	}
}

func (u *Updater) run() {
	os.MkdirAll(u.dir, 0777)
	if u.wantUpdate() {
		l := exec.Command("logger", "-thk")
		c := exec.Command("hk", "update")
		if w, err := l.StdinPipe(); err == nil && l.Start() == nil {
			c.Stdout = w
			c.Stderr = w
		}
		c.Start()
	}
}

func (u *Updater) wantUpdate() bool {
	path := u.dir + upcktimePath
	if Version == "dev" || readTime(path).After(time.Now()) {
		return false
	}
	wait := 24*time.Hour + randDuration(24*time.Hour)
	return writeTime(path, time.Now().Add(wait))
}

func (u *Updater) update() error {
	filename := "hk"
	if runtime.GOOS == "windows" {
		filename += ".exe"
	}
	path, err := exec.LookPath(filename)
	if err != nil {
		return err
	}
	old, err := os.Open(path)
	if err != nil {
		return err
	}
	err = u.fetchInfo()
	if err != nil {
		return err
	}
	if u.info.Version == Version {
		return nil
	}
	bin, err := u.fetchAndApplyPatch(old)
	if err != nil {
		bin, err = u.fetchBin()
		if err != nil {
			return err
		}
	}
	h := sha256.New()
	h.Write(bin)
	if !bytes.Equal(h.Sum(nil), u.info.Sha256) {
		return errors.New("new file hash mismatch after patch")
	}
	return install(old.Name(), bin)
}

func (u *Updater) fetchInfo() error {
	r, err := fetch(u.hkURL + "hk-current-" + plat + ".json")
	if err != nil {
		return err
	}
	defer r.Close()
	err = json.NewDecoder(r).Decode(&u.info)
	if err != nil {
		return err
	}
	if len(u.info.Sha256) != sha256.Size {
		return errors.New("bad cmd hash in info")
	}
	return nil
}

func (u *Updater) fetchAndApplyPatch(old io.Reader) ([]byte, error) {
	r, err := fetch(u.diffURL + slug(Version) + "-to-" + u.info.Version)
	if err != nil {
		return nil, err
	}
	defer r.Close()
	var buf bytes.Buffer
	err = binarydist.Patch(old, &buf, r)
	return buf.Bytes(), err
}

func (u *Updater) fetchBin() ([]byte, error) {
	r, err := fetch(u.binURL + slug(u.info.Version) + ".gz")
	if err != nil {
		return nil, err
	}
	defer r.Close()
	buf := new(bytes.Buffer)
	gz, err := gzip.NewReader(r)
	if err != nil {
		return nil, err
	}
	if _, err = io.Copy(buf, gz); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

func install(name string, p []byte) error {
	execDir := filepath.Dir(name)
	part := filepath.Join(execDir, "hk.part")
	err := ioutil.WriteFile(part, p, 0755)
	if err != nil {
		return err
	}
	defer os.Remove(part)

	// move the existing executable to a new file in the same directory
	oldExecPath := name + ".old"
	err = os.Rename(name, oldExecPath)
	if err != nil {
		return err
	}

	// move the new executable in to become the new program
	err = os.Rename(part, name)

	if err != nil {
		// copy unsuccessful
		_ := os.Rename(oldExecPath, name)
		return err
	} else {
		// copy successful, remove the old binary
		_ = os.Remove(oldExecPath)
	}

	return nil
}

// returns a random duration in [0,n).
func randDuration(n time.Duration) time.Duration {
	return time.Duration(rand.Int63n(int64(n)))
}

func slug(ver string) string {
	return "hk-" + ver + "-" + plat
}

func fetch(url string) (io.ReadCloser, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("bad http status from %s: %v", url, resp.Status)
	}
	return resp.Body, nil
}

func readTime(path string) time.Time {
	p, err := ioutil.ReadFile(path)
	if os.IsNotExist(err) {
		return time.Time{}
	}
	if err != nil {
		return time.Now().Add(1000 * time.Hour)
	}
	t, err := time.Parse(time.RFC3339, string(p))
	if err != nil {
		return time.Now().Add(1000 * time.Hour)
	}
	return t
}

func writeTime(path string, t time.Time) bool {
	return ioutil.WriteFile(path, []byte(t.Format(time.RFC3339)), 0644) == nil
}
<commit_msg>fix an issue where hk update would fail on windows + go fmt<commit_after>package main

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/kr/binarydist"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"
)

const (
	upcktimePath = "cktime"
	plat         = runtime.GOOS + "-" + runtime.GOARCH
)

const devValidTime = 7 * 24 * time.Hour

// Update protocol.
//
//   GET hk.heroku.com/hk-current-linux-amd64.json
//
//   200 ok
//   {
//       "Version": "2",
//       "Sha256": "..." // base64
//   }
//
// then
//
//   GET hkpatch.s3.amazonaws.com/hk-1-linux-amd64-to-2
//
//   200 ok
//   [bsdiff data]
//
// or
//
//   GET hkdist.s3.amazonaws.com/hk-2-linux-amd64.gz
//
//   200 ok
//   [gzipped executable data]
type Updater struct {
	hkURL   string
	binURL  string
	diffURL string
	dir     string
	info    struct {
		Version string
		Sha256  []byte
	}
}

func (u *Updater) run() {
	os.MkdirAll(u.dir, 0777)
	if u.wantUpdate() {
		l := exec.Command("logger", "-thk")
		c := exec.Command("hk", "update")
		if w, err := l.StdinPipe(); err == nil && l.Start() == nil {
			c.Stdout = w
			c.Stderr = w
		}
		c.Start()
	}
}

func (u *Updater) wantUpdate() bool {
	path := u.dir + upcktimePath
	if Version == "dev" || readTime(path).After(time.Now()) {
		return false
	}
	wait := 24*time.Hour + randDuration(24*time.Hour)
	return writeTime(path, time.Now().Add(wait))
}

func (u *Updater) update() error {
	filename := "hk"
	if runtime.GOOS == "windows" {
		filename += ".exe"
	}
	path, err := exec.LookPath(filename)
	if err != nil {
		return err
	}
	old, err := os.Open(path)
	if err != nil {
		return err
	}
	defer old.Close()
	err = u.fetchInfo()
	if err != nil {
		return err
	}
	if u.info.Version == Version {
		return nil
	}
	bin, err := u.fetchAndApplyPatch(old)
	if err != nil {
		bin, err = u.fetchBin()
		if err != nil {
			return err
		}
	}
	h := sha256.New()
	h.Write(bin)
	if !bytes.Equal(h.Sum(nil), u.info.Sha256) {
		return errors.New("new file hash mismatch after patch")
	}
	oldName := old.Name()
	// close the old binary before installing because on windows
	// it can't be renamed if a handle to the file is still open
	old.Close()
	return install(oldName, bin)
}

func (u *Updater) fetchInfo() error {
	r, err := fetch(u.hkURL + "hk-current-" + plat + ".json")
	if err != nil {
		return err
	}
	defer r.Close()
	err = json.NewDecoder(r).Decode(&u.info)
	if err != nil {
		return err
	}
	if len(u.info.Sha256) != sha256.Size {
		return errors.New("bad cmd hash in info")
	}
	return nil
}

func (u *Updater) fetchAndApplyPatch(old io.Reader) ([]byte, error) {
	r, err := fetch(u.diffURL + slug(Version) + "-to-" + u.info.Version)
	if err != nil {
		return nil, err
	}
	defer r.Close()
	var buf bytes.Buffer
	err = binarydist.Patch(old, &buf, r)
	return buf.Bytes(), err
}

func (u *Updater) fetchBin() ([]byte, error) {
	r, err := fetch(u.binURL + slug(u.info.Version) + ".gz")
	if err != nil {
		return nil, err
	}
	defer r.Close()
	buf := new(bytes.Buffer)
	gz, err := gzip.NewReader(r)
	if err != nil {
		return nil, err
	}
	if _, err = io.Copy(buf, gz); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

func install(name string, p []byte) error {
	execDir := filepath.Dir(name)
	part := filepath.Join(execDir, "hk.part")
	err := ioutil.WriteFile(part, p, 0755)
	if err != nil {
		return err
	}
	defer os.Remove(part)

	// move the existing executable to a new file in the same directory
	oldExecPath := name + ".old"
	err = os.Rename(name, oldExecPath)
	if err != nil {
		return err
	}

	// move the new executable in to become the new program
	err = os.Rename(part, name)

	if err != nil {
		// copy unsuccessful
		_ = os.Rename(oldExecPath, name)
		return err
	} else {
		// copy successful, remove the old binary
		_ = os.Remove(oldExecPath)
	}

	return nil
}

// returns a random duration in [0,n).
func randDuration(n time.Duration) time.Duration {
	return time.Duration(rand.Int63n(int64(n)))
}

func slug(ver string) string {
	return "hk-" + ver + "-" + plat
}

func fetch(url string) (io.ReadCloser, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("bad http status from %s: %v", url, resp.Status)
	}
	return resp.Body, nil
}

func readTime(path string) time.Time {
	p, err := ioutil.ReadFile(path)
	if os.IsNotExist(err) {
		return time.Time{}
	}
	if err != nil {
		return time.Now().Add(1000 * time.Hour)
	}
	t, err := time.Parse(time.RFC3339, string(p))
	if err != nil {
		return time.Now().Add(1000 * time.Hour)
	}
	return t
}

func writeTime(path string, t time.Time) bool {
	return ioutil.WriteFile(path, []byte(t.Format(time.RFC3339)), 0644) == nil
}
<|endoftext|>


<commit_before>package geoloader

import (
	"context"
	"log"
	"regexp"
	"sort"
	"sync"
	"time"

	"cloud.google.com/go/storage"
	"google.golang.org/api/iterator"

	"github.com/m-lab/annotation-service/api"
)

// GeoLite2StartDate is the date we have the first GeoLite2 dataset.
// Any request earlier than this date using legacy binary datasets
// later than this date using GeoLite2 datasets
// TODO make this local
var GeoLite2StartDate = time.Unix(1502755200, 0) //"August 15, 2017"

// earliestArchiveDate is the date of the earliest archived dataset.
var earliestArchiveDate = time.Unix(1377648000, 0) // "August 28, 2013")

// datasetDir stores info on all the available datasets.  It is initially empty, just to
// provide the LatestDate() function.
// The current directory is regarded as immutable, but the pointer is dynamically updated, so accesses
// should only be done through GetDirectory().
var datasetDir = &directory{}
var mutex sync.RWMutex // lock to be held when accessing or updating datasetDir pointer.

type dateEntry struct {
	date      time.Time
	filenames []string
}

// directory maintains a list of datasets.
type directory struct {
	entries map[string]*dateEntry // Map to filenames associated with date.
	dates   []string              // Date strings associated with files.
}

func newDirectory(size int) directory {
	return directory{entries: make(map[string]*dateEntry, size), dates: make([]string, 0, size)}
}

// Insert inserts a new filename into the directory at the given date.
// NOTE: This does not detect or eliminate duplicates.
func (dir *directory) insert(date time.Time, fn string) {
	dateString := date.Format("20060102")
	entry, ok := dir.entries[dateString]
	if !ok {
		// Insert the new date into the date slice.
		index := sort.SearchStrings(dir.dates, dateString)
		dir.dates = append(dir.dates, "")
		copy(dir.dates[index+1:], dir.dates[index:])
		dir.dates[index] = dateString

		// Create new entry for the date.
		entry = &dateEntry{filenames: make([]string, 0, 3), date: date}
		dir.entries[dateString] = entry
	}

	entry.filenames = append(entry.filenames, fn)
}

func (dir *directory) lastDate() time.Time {
	if len(dir.dates) < 1 {
		return time.Time{}
	}
	d := dir.dates[len(dir.dates)-1]
	return dir.entries[d].date
}

// LastBefore returns the filename associated with the provided date.
func (dir *directory) LastBefore(date time.Time) string {
	if len(dir.dates) == 0 {
		return ""
	}

	// Add one day, so that we use dataset for same date published
	// TODO perhaps we should use the previous dataset???
	date = date.Add(24 * time.Hour)

	dateString := date.Format("20060102")
	index := sort.SearchStrings(dir.dates, dateString)
	if index == 0 {
		return dir.entries[dir.dates[index]].filenames[0]
	}
	return dir.entries[dir.dates[index-1]].filenames[0]
}

// TODO: These regex are duplicated in geolite2 and legacy packages.
// This is the regex used to filter for which files we want to consider acceptable for using with Geolite2
var GeoLite2Regex = regexp.MustCompile(`Maxmind/\d{4}/\d{2}/\d{2}/\d{8}T\d{6}Z-GeoLite2-City-CSV\.zip`)

// This is the regex used to filter for which files we want to consider acceptable for using with legacy dataset
var GeoLegacyRegex = regexp.MustCompile(`.*-GeoLiteCity.dat.*`)
var GeoLegacyv6Regex = regexp.MustCompile(`.*-GeoLiteCityv6.dat.*`)

// UpdateArchivedFilenames extracts the dataset filenames from downloader bucket
// This job is run at the beginning of deployment and daily cron job.
func UpdateArchivedFilenames() error {
	old := getDirectory()
	size := len(old.dates) + 2
	dir := directory{entries: make(map[string]*dateEntry, size), dates: make([]string, 0, size)}

	ctx := context.Background()
	client, err := storage.NewClient(ctx)
	if err != nil {
		return err
	}
	prospectiveFiles := client.Bucket(api.MaxmindBucketName).Objects(ctx, &storage.Query{Prefix: api.MaxmindPrefix})
	for file, err := prospectiveFiles.Next(); err != iterator.Done; file, err = prospectiveFiles.Next() {
		if err != nil {
			return err
		}
		if !GeoLite2Regex.MatchString(file.Name) && !GeoLegacyRegex.MatchString(file.Name) && !GeoLegacyv6Regex.MatchString(file.Name) {
			continue
		}
		// We archived but do not use legacy datasets after GeoLite2StartDate.
		fileDate, err := api.ExtractDateFromFilename(file.Name)
		if err != nil {
			continue
		}
		if !fileDate.Before(GeoLite2StartDate) && !GeoLite2Regex.MatchString(file.Name) {
			continue
		}

		dir.insert(fileDate, file.Name)
	}
	if err != nil {
		log.Println(err)
	}

	mutex.Lock()
	datasetDir = &dir
	mutex.Unlock()

	return nil
}

func getDirectory() *directory {
	mutex.RLock()
	defer mutex.RUnlock()
	return datasetDir
}

// Latest returns the date of the latest dataset.
// May return time.Time{} if no dates have been loaded.
func Latest() time.Time {
	dd := getDirectory()
	return dd.lastDate()
}

// LastBefore returns the dataset filename for annotating the requested date.
func LastBefore(date time.Time) string {
	dd := getDirectory()
	return dd.LastBefore(date)
}
<commit_msg>export Insert for testing<commit_after>package geoloader

import (
	"context"
	"log"
	"regexp"
	"sort"
	"sync"
	"time"

	"cloud.google.com/go/storage"
	"google.golang.org/api/iterator"

	"github.com/m-lab/annotation-service/api"
)

// GeoLite2StartDate is the date we have the first GeoLite2 dataset.
// Any request earlier than this date using legacy binary datasets
// later than this date using GeoLite2 datasets
// TODO make this local
var GeoLite2StartDate = time.Unix(1502755200, 0) //"August 15, 2017"

// earliestArchiveDate is the date of the earliest archived dataset.
var earliestArchiveDate = time.Unix(1377648000, 0) // "August 28, 2013")

// datasetDir stores info on all the available datasets.  It is initially empty, just to
// provide the LatestDate() function.
// The current directory is regarded as immutable, but the pointer is dynamically updated, so accesses
// should only be done through GetDirectory().
var datasetDir = &directory{}
var mutex sync.RWMutex // lock to be held when accessing or updating datasetDir pointer.

type dateEntry struct {
	date      time.Time
	filenames []string
}

// directory maintains a list of datasets.
type directory struct {
	entries map[string]*dateEntry // Map to filenames associated with date.
	dates   []string              // Date strings associated with files.
}

func newDirectory(size int) directory {
	return directory{entries: make(map[string]*dateEntry, size), dates: make([]string, 0, size)}
}

// Insert inserts a new filename into the directory at the given date.
// NOTE: This does not detect or eliminate duplicates.
// TODO - make this local.
func (dir *directory) Insert(date time.Time, fn string) {
	dateString := date.Format("20060102")
	entry, ok := dir.entries[dateString]
	if !ok {
		// Insert the new date into the date slice.
		index := sort.SearchStrings(dir.dates, dateString)
		dir.dates = append(dir.dates, "")
		copy(dir.dates[index+1:], dir.dates[index:])
		dir.dates[index] = dateString

		// Create new entry for the date.
		entry = &dateEntry{filenames: make([]string, 0, 3), date: date}
		dir.entries[dateString] = entry
	}

	entry.filenames = append(entry.filenames, fn)
}

func (dir *directory) lastDate() time.Time {
	if len(dir.dates) < 1 {
		return time.Time{}
	}
	d := dir.dates[len(dir.dates)-1]
	return dir.entries[d].date
}

// LastBefore returns the filename associated with the provided date.
func (dir *directory) LastBefore(date time.Time) string {
	if len(dir.dates) == 0 {
		return ""
	}

	// Add one day, so that we use dataset for same date published
	// TODO perhaps we should use the previous dataset???
	date = date.Add(24 * time.Hour)

	dateString := date.Format("20060102")
	index := sort.SearchStrings(dir.dates, dateString)
	if index == 0 {
		return dir.entries[dir.dates[index]].filenames[0]
	}
	return dir.entries[dir.dates[index-1]].filenames[0]
}

// TODO: These regex are duplicated in geolite2 and legacy packages.
// This is the regex used to filter for which files we want to consider acceptable for using with Geolite2
var GeoLite2Regex = regexp.MustCompile(`Maxmind/\d{4}/\d{2}/\d{2}/\d{8}T\d{6}Z-GeoLite2-City-CSV\.zip`)

// This is the regex used to filter for which files we want to consider acceptable for using with legacy dataset
var GeoLegacyRegex = regexp.MustCompile(`.*-GeoLiteCity.dat.*`)
var GeoLegacyv6Regex = regexp.MustCompile(`.*-GeoLiteCityv6.dat.*`)

// UpdateArchivedFilenames extracts the dataset filenames from downloader bucket
// This job is run at the beginning of deployment and daily cron job.
func UpdateArchivedFilenames() error {
	old := getDirectory()
	size := len(old.dates) + 2
	dir := directory{entries: make(map[string]*dateEntry, size), dates: make([]string, 0, size)}

	ctx := context.Background()
	client, err := storage.NewClient(ctx)
	if err != nil {
		return err
	}
	prospectiveFiles := client.Bucket(api.MaxmindBucketName).Objects(ctx, &storage.Query{Prefix: api.MaxmindPrefix})
	for file, err := prospectiveFiles.Next(); err != iterator.Done; file, err = prospectiveFiles.Next() {
		if err != nil {
			return err
		}
		if !GeoLite2Regex.MatchString(file.Name) && !GeoLegacyRegex.MatchString(file.Name) && !GeoLegacyv6Regex.MatchString(file.Name) {
			continue
		}
		// We archived but do not use legacy datasets after GeoLite2StartDate.
		fileDate, err := api.ExtractDateFromFilename(file.Name)
		if err != nil {
			continue
		}
		if !fileDate.Before(GeoLite2StartDate) && !GeoLite2Regex.MatchString(file.Name) {
			continue
		}

		dir.Insert(fileDate, file.Name)
	}
	if err != nil {
		log.Println(err)
	}

	mutex.Lock()
	datasetDir = &dir
	mutex.Unlock()

	return nil
}

func getDirectory() *directory {
	mutex.RLock()
	defer mutex.RUnlock()
	return datasetDir
}

// Latest returns the date of the latest dataset.
// May return time.Time{} if no dates have been loaded.
func Latest() time.Time {
	dd := getDirectory()
	return dd.lastDate()
}

// LastBefore returns the dataset filename for annotating the requested date.
func LastBefore(date time.Time) string {
	dd := getDirectory()
	return dd.LastBefore(date)
}
<|endoftext|>


<commit_before>package azure

import (
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/svanharmelen/azure-sdk-for-go/management"
	"github.com/svanharmelen/azure-sdk-for-go/management/virtualmachinedisk"
)

const dataDiskBlobStorageURL = "http://%s.blob.core.windows.net/disks/%s.vhd"

func resourceAzureDataDisk() *schema.Resource {
	return &schema.Resource{
		Create: resourceAzureDataDiskCreate,
		Read:   resourceAzureDataDiskRead,
		Update: resourceAzureDataDiskUpdate,
		Delete: resourceAzureDataDiskDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
			},

			"label": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
			},

			"lun": &schema.Schema{
				Type:     schema.TypeInt,
				Required: true,
			},

			"size": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
			},

			"caching": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Default:  "None",
			},

			"storage": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"media_link": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
			},

			"source_media_link": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"virtual_machine": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
		},
	}
}

func resourceAzureDataDiskCreate(d *schema.ResourceData, meta interface{}) error {
	mc := meta.(*Client).mgmtClient

	if err := verifyDataDiskParameters(d); err != nil {
		return err
	}

	lun := d.Get("lun").(int)
	vm := d.Get("virtual_machine").(string)

	label := d.Get("label").(string)
	if label == "" {
		label = fmt.Sprintf("%s-%d", vm, lun)
	}

	p := virtualmachinedisk.CreateDataDiskParameters{
		DiskLabel:           label,
		Lun:                 lun,
		LogicalDiskSizeInGB: d.Get("size").(int),
		HostCaching:         hostCaching(d),
		MediaLink:           mediaLink(d),
		SourceMediaLink:     d.Get("source_media_link").(string),
	}

	if name, ok := d.GetOk("name"); ok {
		p.DiskName = name.(string)
	}

	log.Printf("[DEBUG] Adding data disk %d to instance: %s", lun, vm)
	req, err := virtualmachinedisk.NewClient(mc).AddDataDisk(vm, vm, vm, p)
	if err != nil {
		return fmt.Errorf("Error adding data disk %d to instance %s: %s", lun, vm, err)
	}

	// Wait until the data disk is added
	if err := mc.WaitForOperation(req, nil); err != nil {
		return fmt.Errorf(
			"Error waiting for data disk %d to be added to instance %s: %s", lun, vm, err)
	}

	log.Printf("[DEBUG] Retrieving data disk %d from instance %s", lun, vm)
	disk, err := virtualmachinedisk.NewClient(mc).GetDataDisk(vm, vm, vm, lun)
	if err != nil {
		return fmt.Errorf("Error retrieving data disk %d from instance %s: %s", lun, vm, err)
	}

	d.SetId(disk.DiskName)

	return resourceAzureDataDiskRead(d, meta)
}

func resourceAzureDataDiskRead(d *schema.ResourceData, meta interface{}) error {
	mc := meta.(*Client).mgmtClient

	lun := d.Get("lun").(int)
	vm := d.Get("virtual_machine").(string)

	log.Printf("[DEBUG] Retrieving data disk: %s", d.Id())
	datadisk, err := virtualmachinedisk.NewClient(mc).GetDataDisk(vm, vm, vm, lun)
	if err != nil {
		if management.IsResourceNotFoundError(err) {
			d.SetId("")
			return nil
		}
		return fmt.Errorf("Error retrieving data disk %s: %s", d.Id(), err)
	}

	d.Set("name", datadisk.DiskName)
	d.Set("label", datadisk.DiskLabel)
	d.Set("lun", datadisk.Lun)
	d.Set("size", datadisk.LogicalDiskSizeInGB)
	d.Set("caching", datadisk.HostCaching)
	d.Set("media_link", datadisk.MediaLink)

	log.Printf("[DEBUG] Retrieving disk: %s", d.Id())
	disk, err := virtualmachinedisk.NewClient(mc).GetDisk(d.Id())
	if err != nil {
		return fmt.Errorf("Error retrieving disk %s: %s", d.Id(), err)
	}

	d.Set("virtual_machine", disk.AttachedTo.RoleName)

	return nil
}

func resourceAzureDataDiskUpdate(d *schema.ResourceData, meta interface{}) error {
	mc := meta.(*Client).mgmtClient
	diskClient := virtualmachinedisk.NewClient(mc)

	lun := d.Get("lun").(int)
	vm := d.Get("virtual_machine").(string)

	if d.HasChange("lun") || d.HasChange("size") || d.HasChange("virtual_machine") {
		olun, _ := d.GetChange("lun")
		ovm, _ := d.GetChange("virtual_machine")

		log.Printf("[DEBUG] Detaching data disk: %s", d.Id())
		req, err := diskClient.
			DeleteDataDisk(ovm.(string), ovm.(string), ovm.(string), olun.(int), false)
		if err != nil {
			return fmt.Errorf("Error detaching data disk %s: %s", d.Id(), err)
		}

		// Wait until the data disk is detached
		if err := mc.WaitForOperation(req, nil); err != nil {
			return fmt.Errorf(
				"Error waiting for data disk %s to be detached: %s", d.Id(), err)
		}

		log.Printf("[DEBUG] Verifying data disk %s is properly detached...", d.Id())
		for i := 0; i < 6; i++ {
			disk, err := diskClient.GetDisk(d.Id())
			if err != nil {
				return fmt.Errorf("Error retrieving disk %s: %s", d.Id(), err)
			}

			// Check if the disk is really detached
			if disk.AttachedTo.RoleName == "" {
				break
			}

			// If not, wait 30 seconds and try it again...
			time.Sleep(time.Duration(30 * time.Second))
		}

		if d.HasChange("size") {
			p := virtualmachinedisk.UpdateDiskParameters{
				DiskName:        d.Id(),
				Label:           d.Get("label").(string),
				ResizedSizeInGB: d.Get("size").(int),
			}

			log.Printf("[DEBUG] Updating disk: %s", d.Id())
			req, err := diskClient.UpdateDisk(d.Id(), p)
			if err != nil {
				return fmt.Errorf("Error updating disk %s: %s", d.Id(), err)
			}

			// Wait until the disk is updated
			if err := mc.WaitForOperation(req, nil); err != nil {
				return fmt.Errorf(
					"Error waiting for disk %s to be updated: %s", d.Id(), err)
			}
		}

		p := virtualmachinedisk.CreateDataDiskParameters{
			DiskName:    d.Id(),
			Lun:         lun,
			HostCaching: hostCaching(d),
			MediaLink:   mediaLink(d),
		}

		log.Printf("[DEBUG] Attaching data disk: %s", d.Id())
		req, err = diskClient.AddDataDisk(vm, vm, vm, p)
		if err != nil {
			return fmt.Errorf("Error attaching data disk %s to instance %s: %s", d.Id(), vm, err)
		}

		// Wait until the data disk is attached
		if err := mc.WaitForOperation(req, nil); err != nil {
			return fmt.Errorf(
				"Error waiting for data disk %s to be attached to instance %s: %s", d.Id(), vm, err)
		}

		// Make sure we return here since all possible changes are
		// already updated if we reach this point
		return nil
	}

	if d.HasChange("caching") {
		p := virtualmachinedisk.UpdateDataDiskParameters{
			DiskName:    d.Id(),
			Lun:         lun,
			HostCaching: hostCaching(d),
			MediaLink:   mediaLink(d),
		}

		log.Printf("[DEBUG] Updating data disk: %s", d.Id())
		req, err := diskClient.UpdateDataDisk(vm, vm, vm, lun, p)
		if err != nil {
			return fmt.Errorf("Error updating data disk %s: %s", d.Id(), err)
		}

		// Wait until the data disk is updated
		if err := mc.WaitForOperation(req, nil); err != nil {
			return fmt.Errorf(
				"Error waiting for data disk %s to be updated: %s", d.Id(), err)
		}
	}

	return resourceAzureDataDiskRead(d, meta)
}

func resourceAzureDataDiskDelete(d *schema.ResourceData, meta interface{}) error {
	mc := meta.(*Client).mgmtClient

	lun := d.Get("lun").(int)
	vm := d.Get("virtual_machine").(string)

	// If a name was not supplied, it means we created a new emtpy disk and we now want to
	// delete that disk again. Otherwise we only want to detach the disk and keep the blob.
	_, removeBlob := d.GetOk("name")

	log.Printf("[DEBUG] Detaching data disk %s with removeBlob = %t", d.Id(), removeBlob)
	req, err := virtualmachinedisk.NewClient(mc).DeleteDataDisk(vm, vm, vm, lun, removeBlob)
	if err != nil {
		return fmt.Errorf(
			"Error detaching data disk %s with removeBlob = %t: %s", d.Id(), removeBlob, err)
	}

	// Wait until the data disk is detached and optionally deleted
	if err := mc.WaitForOperation(req, nil); err != nil {
		return fmt.Errorf(
			"Error waiting for data disk %s to be detached with removeBlob = %t: %s",
			d.Id(), removeBlob, err)
	}

	d.SetId("")

	return nil
}

func hostCaching(d *schema.ResourceData) virtualmachinedisk.HostCachingType {
	switch d.Get("caching").(string) {
	case "ReadOnly":
		return virtualmachinedisk.HostCachingTypeReadOnly
	case "ReadWrite":
		return virtualmachinedisk.HostCachingTypeReadWrite
	default:
		return virtualmachinedisk.HostCachingTypeNone
	}
}

func mediaLink(d *schema.ResourceData) string {
	mediaLink, ok := d.GetOk("media_link")
	if ok {
		return mediaLink.(string)
	}

	name, ok := d.GetOk("name")
	if !ok {
		name = fmt.Sprintf("%s-%d", d.Get("virtual_machine").(string), d.Get("lun").(int))
	}

	return fmt.Sprintf(dataDiskBlobStorageURL, d.Get("storage").(string), name.(string))
}

func verifyDataDiskParameters(d *schema.ResourceData) error {
	caching := d.Get("caching").(string)
	if caching != "None" && caching != "ReadOnly" && caching != "ReadWrite" {
		return fmt.Errorf(
			"Invalid caching type %s! Valid options are 'None', 'ReadOnly' and 'ReadWrite'.", caching)
	}

	if _, ok := d.GetOk("media_link"); !ok {
		if _, ok := d.GetOk("storage"); !ok {
			return fmt.Errorf("If not supplying 'media_link', you must supply 'storage'.")
		}
	}

	return nil
}
<commit_msg>Very minor change needed to realign with the updated Azure SDK<commit_after>package azure

import (
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/svanharmelen/azure-sdk-for-go/management"
	"github.com/svanharmelen/azure-sdk-for-go/management/virtualmachinedisk"
)

const dataDiskBlobStorageURL = "http://%s.blob.core.windows.net/disks/%s.vhd"

func resourceAzureDataDisk() *schema.Resource {
	return &schema.Resource{
		Create: resourceAzureDataDiskCreate,
		Read:   resourceAzureDataDiskRead,
		Update: resourceAzureDataDiskUpdate,
		Delete: resourceAzureDataDiskDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
			},

			"label": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
			},

			"lun": &schema.Schema{
				Type:     schema.TypeInt,
				Required: true,
			},

			"size": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
			},

			"caching": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Default:  "None",
			},

			"storage": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"media_link": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
			},

			"source_media_link": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"virtual_machine": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
		},
	}
}

func resourceAzureDataDiskCreate(d *schema.ResourceData, meta interface{}) error {
	mc := meta.(*Client).mgmtClient

	if err := verifyDataDiskParameters(d); err != nil {
		return err
	}

	lun := d.Get("lun").(int)
	vm := d.Get("virtual_machine").(string)

	label := d.Get("label").(string)
	if label == "" {
		label = fmt.Sprintf("%s-%d", vm, lun)
	}

	p := virtualmachinedisk.CreateDataDiskParameters{
		DiskLabel:           label,
		Lun:                 lun,
		LogicalDiskSizeInGB: d.Get("size").(int),
		HostCaching:         hostCaching(d),
		MediaLink:           mediaLink(d),
		SourceMediaLink:     d.Get("source_media_link").(string),
	}

	if name, ok := d.GetOk("name"); ok {
		p.DiskName = name.(string)
	}

	log.Printf("[DEBUG] Adding data disk %d to instance: %s", lun, vm)
	req, err := virtualmachinedisk.NewClient(mc).AddDataDisk(vm, vm, vm, p)
	if err != nil {
		return fmt.Errorf("Error adding data disk %d to instance %s: %s", lun, vm, err)
	}

	// Wait until the data disk is added
	if err := mc.WaitForOperation(req, nil); err != nil {
		return fmt.Errorf(
			"Error waiting for data disk %d to be added to instance %s: %s", lun, vm, err)
	}

	log.Printf("[DEBUG] Retrieving data disk %d from instance %s", lun, vm)
	disk, err := virtualmachinedisk.NewClient(mc).GetDataDisk(vm, vm, vm, lun)
	if err != nil {
		return fmt.Errorf("Error retrieving data disk %d from instance %s: %s", lun, vm, err)
	}

	d.SetId(disk.DiskName)

	return resourceAzureDataDiskRead(d, meta)
}

func resourceAzureDataDiskRead(d *schema.ResourceData, meta interface{}) error {
	mc := meta.(*Client).mgmtClient

	lun := d.Get("lun").(int)
	vm := d.Get("virtual_machine").(string)

	log.Printf("[DEBUG] Retrieving data disk: %s", d.Id())
	datadisk, err := virtualmachinedisk.NewClient(mc).GetDataDisk(vm, vm, vm, lun)
	if err != nil {
		if management.IsResourceNotFoundError(err) {
			d.SetId("")
			return nil
		}
		return fmt.Errorf("Error retrieving data disk %s: %s", d.Id(), err)
	}

	d.Set("name", datadisk.DiskName)
	d.Set("label", datadisk.DiskLabel)
	d.Set("lun", datadisk.Lun)
	d.Set("size", datadisk.LogicalDiskSizeInGB)
	d.Set("caching", datadisk.HostCaching)
	d.Set("media_link", datadisk.MediaLink)

	log.Printf("[DEBUG] Retrieving disk: %s", d.Id())
	disk, err := virtualmachinedisk.NewClient(mc).GetDisk(d.Id())
	if err != nil {
		return fmt.Errorf("Error retrieving disk %s: %s", d.Id(), err)
	}

	d.Set("virtual_machine", disk.AttachedTo.RoleName)

	return nil
}

func resourceAzureDataDiskUpdate(d *schema.ResourceData, meta interface{}) error {
	mc := meta.(*Client).mgmtClient
	diskClient := virtualmachinedisk.NewClient(mc)

	lun := d.Get("lun").(int)
	vm := d.Get("virtual_machine").(string)

	if d.HasChange("lun") || d.HasChange("size") || d.HasChange("virtual_machine") {
		olun, _ := d.GetChange("lun")
		ovm, _ := d.GetChange("virtual_machine")

		log.Printf("[DEBUG] Detaching data disk: %s", d.Id())
		req, err := diskClient.
			DeleteDataDisk(ovm.(string), ovm.(string), ovm.(string), olun.(int), false)
		if err != nil {
			return fmt.Errorf("Error detaching data disk %s: %s", d.Id(), err)
		}

		// Wait until the data disk is detached
		if err := mc.WaitForOperation(req, nil); err != nil {
			return fmt.Errorf(
				"Error waiting for data disk %s to be detached: %s", d.Id(), err)
		}

		log.Printf("[DEBUG] Verifying data disk %s is properly detached...", d.Id())
		for i := 0; i < 6; i++ {
			disk, err := diskClient.GetDisk(d.Id())
			if err != nil {
				return fmt.Errorf("Error retrieving disk %s: %s", d.Id(), err)
			}

			// Check if the disk is really detached
			if disk.AttachedTo.RoleName == "" {
				break
			}

			// If not, wait 30 seconds and try it again...
			time.Sleep(time.Duration(30 * time.Second))
		}

		if d.HasChange("size") {
			p := virtualmachinedisk.UpdateDiskParameters{
				Name:            d.Id(),
				Label:           d.Get("label").(string),
				ResizedSizeInGB: d.Get("size").(int),
			}

			log.Printf("[DEBUG] Updating disk: %s", d.Id())
			req, err := diskClient.UpdateDisk(d.Id(), p)
			if err != nil {
				return fmt.Errorf("Error updating disk %s: %s", d.Id(), err)
			}

			// Wait until the disk is updated
			if err := mc.WaitForOperation(req, nil); err != nil {
				return fmt.Errorf(
					"Error waiting for disk %s to be updated: %s", d.Id(), err)
			}
		}

		p := virtualmachinedisk.CreateDataDiskParameters{
			DiskName:    d.Id(),
			Lun:         lun,
			HostCaching: hostCaching(d),
			MediaLink:   mediaLink(d),
		}

		log.Printf("[DEBUG] Attaching data disk: %s", d.Id())
		req, err = diskClient.AddDataDisk(vm, vm, vm, p)
		if err != nil {
			return fmt.Errorf("Error attaching data disk %s to instance %s: %s", d.Id(), vm, err)
		}

		// Wait until the data disk is attached
		if err := mc.WaitForOperation(req, nil); err != nil {
			return fmt.Errorf(
				"Error waiting for data disk %s to be attached to instance %s: %s", d.Id(), vm, err)
		}

		// Make sure we return here since all possible changes are
		// already updated if we reach this point
		return nil
	}

	if d.HasChange("caching") {
		p := virtualmachinedisk.UpdateDataDiskParameters{
			DiskName:    d.Id(),
			Lun:         lun,
			HostCaching: hostCaching(d),
			MediaLink:   mediaLink(d),
		}

		log.Printf("[DEBUG] Updating data disk: %s", d.Id())
		req, err := diskClient.UpdateDataDisk(vm, vm, vm, lun, p)
		if err != nil {
			return fmt.Errorf("Error updating data disk %s: %s", d.Id(), err)
		}

		// Wait until the data disk is updated
		if err := mc.WaitForOperation(req, nil); err != nil {
			return fmt.Errorf(
				"Error waiting for data disk %s to be updated: %s", d.Id(), err)
		}
	}

	return resourceAzureDataDiskRead(d, meta)
}

func resourceAzureDataDiskDelete(d *schema.ResourceData, meta interface{}) error {
	mc := meta.(*Client).mgmtClient

	lun := d.Get("lun").(int)
	vm := d.Get("virtual_machine").(string)

	// If a name was not supplied, it means we created a new emtpy disk and we now want to
	// delete that disk again. Otherwise we only want to detach the disk and keep the blob.
	_, removeBlob := d.GetOk("name")

	log.Printf("[DEBUG] Detaching data disk %s with removeBlob = %t", d.Id(), removeBlob)
	req, err := virtualmachinedisk.NewClient(mc).DeleteDataDisk(vm, vm, vm, lun, removeBlob)
	if err != nil {
		return fmt.Errorf(
			"Error detaching data disk %s with removeBlob = %t: %s", d.Id(), removeBlob, err)
	}

	// Wait until the data disk is detached and optionally deleted
	if err := mc.WaitForOperation(req, nil); err != nil {
		return fmt.Errorf(
			"Error waiting for data disk %s to be detached with removeBlob = %t: %s",
			d.Id(), removeBlob, err)
	}

	d.SetId("")

	return nil
}

func hostCaching(d *schema.ResourceData) virtualmachinedisk.HostCachingType {
	switch d.Get("caching").(string) {
	case "ReadOnly":
		return virtualmachinedisk.HostCachingTypeReadOnly
	case "ReadWrite":
		return virtualmachinedisk.HostCachingTypeReadWrite
	default:
		return virtualmachinedisk.HostCachingTypeNone
	}
}

func mediaLink(d *schema.ResourceData) string {
	mediaLink, ok := d.GetOk("media_link")
	if ok {
		return mediaLink.(string)
	}

	name, ok := d.GetOk("name")
	if !ok {
		name = fmt.Sprintf("%s-%d", d.Get("virtual_machine").(string), d.Get("lun").(int))
	}

	return fmt.Sprintf(dataDiskBlobStorageURL, d.Get("storage").(string), name.(string))
}

func verifyDataDiskParameters(d *schema.ResourceData) error {
	caching := d.Get("caching").(string)
	if caching != "None" && caching != "ReadOnly" && caching != "ReadWrite" {
		return fmt.Errorf(
			"Invalid caching type %s! Valid options are 'None', 'ReadOnly' and 'ReadWrite'.", caching)
	}

	if _, ok := d.GetOk("media_link"); !ok {
		if _, ok := d.GetOk("storage"); !ok {
			return fmt.Errorf("If not supplying 'media_link', you must supply 'storage'.")
		}
	}

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package gce

import (
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	"gopkg.in/gcfg.v1"

	"github.com/golang/glog"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	gce "google.golang.org/api/compute/v1"
	provider_gce "k8s.io/kubernetes/pkg/cloudprovider/providers/gce"
	"k8s.io/kubernetes/pkg/util/wait"
)

const (
	operationWaitTimeout  = 5 * time.Second
	operationPollInterval = 100 * time.Millisecond
)

type migInformation struct {
	config   *Mig
	basename string
}

// GceManager is handles gce communication and data caching.
type GceManager struct {
	migs     []*migInformation
	migCache map[GceRef]*Mig

	service    *gce.Service
	cacheMutex sync.Mutex
}

// CreateGceManager constructs gceManager object.
func CreateGceManager(configReader io.Reader) (*GceManager, error) {
	// Create Google Compute Engine token.
	tokenSource := google.ComputeTokenSource("")
	if configReader != nil {
		var cfg provider_gce.Config
		if err := gcfg.ReadInto(&cfg, configReader); err != nil {
			glog.Errorf("Couldn't read config: %v", err)
			return nil, err
		}
		if cfg.Global.TokenURL == "" {
			glog.Warning("Empty tokenUrl in cloud config")
		} else {
			glog.Infof("Using TokenSource from config %#v", tokenSource)
			tokenSource = provider_gce.NewAltTokenSource(cfg.Global.TokenURL, cfg.Global.TokenBody)
		}
	} else {
		glog.Infof("Using default TokenSource %#v", tokenSource)
	}

	// Create Google Compute Engine service.
	client := oauth2.NewClient(oauth2.NoContext, tokenSource)
	gceService, err := gce.New(client)
	if err != nil {
		return nil, err
	}
	manager := &GceManager{
		migs:     make([]*migInformation, 0),
		service:  gceService,
		migCache: make(map[GceRef]*Mig),
	}
	go wait.Forever(func() {
		manager.cacheMutex.Lock()
		defer manager.cacheMutex.Unlock()
		if err := manager.regenerateCache(); err != nil {
			glog.Errorf("Error while regenerating Mig cache: %v", err)
		}
	}, time.Hour)
	return manager, nil
}

// RegisterMig registers mig in Gce Manager.
func (m *GceManager) RegisterMig(mig *Mig) {
	m.cacheMutex.Lock()
	defer m.cacheMutex.Unlock()

	m.migs = append(m.migs, &migInformation{
		config: mig,
	})
}

// GetMigSize gets MIG size.
func (m *GceManager) GetMigSize(mig *Mig) (int64, error) {
	igm, err := m.service.InstanceGroupManagers.Get(mig.Project, mig.Zone, mig.Name).Do()
	if err != nil {
		return -1, err
	}
	return igm.TargetSize, nil
}

// SetMigSize sets MIG size.
func (m *GceManager) SetMigSize(mig *Mig, size int64) error {
	op, err := m.service.InstanceGroupManagers.Resize(mig.Project, mig.Zone, mig.Name, size).Do()
	if err != nil {
		return err
	}
	if err := m.waitForOp(op, mig.Project, mig.Zone); err != nil {
		return err
	}
	return nil
}

func (m *GceManager) waitForOp(operation *gce.Operation, project string, zone string) error {
	for start := time.Now(); time.Since(start) < operationWaitTimeout; time.Sleep(operationPollInterval) {
		glog.V(4).Infof("Waiting for operation %s %s %s", project, zone, operation.Name)
		if op, err := m.service.ZoneOperations.Get(project, zone, operation.Name).Do(); err == nil {
			glog.V(4).Infof("Operation %s %s %s status: %s", project, zone, operation.Name, op.Status)
			if op.Status == "DONE" {
				return nil
			}
		} else {
			glog.Warningf("Error while getting operation %s on %s: %v", operation.Name, operation.TargetLink, err)
		}
	}
	return fmt.Errorf("Timeout while waiting for operation %s on %s to complete.", operation.Name, operation.TargetLink)
}

// DeleteInstances deletes the given instances. All instances must be controlled by the same MIG.
func (m *GceManager) DeleteInstances(instances []*GceRef) error {
	if len(instances) == 0 {
		return nil
	}
	commonMig, err := m.GetMigForInstance(instances[0])
	if err != nil {
		return err
	}
	for _, instance := range instances {
		mig, err := m.GetMigForInstance(instance)
		if err != nil {
			return err
		}
		if mig != commonMig {
			return fmt.Errorf("Connot delete instances which don't belong to the same MIG.")
		}
	}

	req := gce.InstanceGroupManagersDeleteInstancesRequest{
		Instances: []string{},
	}
	for _, instance := range instances {
		req.Instances = append(req.Instances, GenerateInstanceUrl(instance.Project, instance.Zone, instance.Name))
	}

	op, err := m.service.InstanceGroupManagers.DeleteInstances(commonMig.Project, commonMig.Zone, commonMig.Name, &req).Do()
	if err != nil {
		return err
	}
	if err := m.waitForOp(op, commonMig.Project, commonMig.Zone); err != nil {
		return err
	}
	return nil
}

// GetMigForInstance returns MigConfig of the given Instance
func (m *GceManager) GetMigForInstance(instance *GceRef) (*Mig, error) {
	m.cacheMutex.Lock()
	defer m.cacheMutex.Unlock()
	if mig, found := m.migCache[*instance]; found {
		return mig, nil
	}

	for _, mig := range m.migs {
		if mig.config.Project == instance.Project &&
			mig.config.Zone == instance.Zone &&
			strings.HasPrefix(instance.Name, mig.basename) {
			if err := m.regenerateCache(); err != nil {
				return nil, fmt.Errorf("Error while looking for MIG for instance %+v, error: %v", *instance, err)
			}
			if mig, found := m.migCache[*instance]; found {
				return mig, nil
			}
			return nil, fmt.Errorf("Instance %+v does not belong to any configured MIG", *instance)
		}
	}
	// Instance doesn't belong to any configured mig.
	return nil, nil
}

func (m *GceManager) regenerateCache() error {
	newMigCache := make(map[GceRef]*Mig)

	for _, migInfo := range m.migs {
		mig := migInfo.config
		glog.V(4).Infof("Regenerating MIG information for %s %s %s", mig.Project, mig.Zone, mig.Name)

		instanceGroupManager, err := m.service.InstanceGroupManagers.Get(mig.Project, mig.Zone, mig.Name).Do()
		if err != nil {
			return err
		}
		migInfo.basename = instanceGroupManager.BaseInstanceName

		instances, err := m.service.InstanceGroupManagers.ListManagedInstances(mig.Project, mig.Zone, mig.Name).Do()
		if err != nil {
			glog.V(4).Infof("Failed MIG info request for %s %s %s: %v", mig.Project, mig.Zone, mig.Name, err)
			return err
		}
		for _, instance := range instances.ManagedInstances {
			project, zone, name, err := ParseInstanceUrl(instance.Instance)
			if err != nil {
				return err
			}
			newMigCache[GceRef{Project: project, Zone: zone, Name: name}] = mig
		}
	}

	m.migCache = newMigCache
	return nil
}

// GetMigNodes returns mig nodes.
func (m *GceManager) GetMigNodes(mig *Mig) ([]string, error) {
	instances, err := m.service.InstanceGroupManagers.ListManagedInstances(mig.Project, mig.Zone, mig.Name).Do()
	if err != nil {
		return []string{}, err
	}
	result := make([]string, 0)
	for _, instance := range instances.ManagedInstances {
		_, _, name, err := ParseInstanceUrl(instance.Instance)
		if err != nil {
			return []string{}, err
		}
		result = append(result, name)
	}
	return result, nil
}
<commit_msg>Cluster-autoscaler: return node names in format of gce://proj/zone/name from GceProvider<commit_after>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package gce

import (
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	"gopkg.in/gcfg.v1"

	"github.com/golang/glog"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	gce "google.golang.org/api/compute/v1"
	provider_gce "k8s.io/kubernetes/pkg/cloudprovider/providers/gce"
	"k8s.io/kubernetes/pkg/util/wait"
)

const (
	operationWaitTimeout  = 5 * time.Second
	operationPollInterval = 100 * time.Millisecond
)

type migInformation struct {
	config   *Mig
	basename string
}

// GceManager is handles gce communication and data caching.
type GceManager struct {
	migs     []*migInformation
	migCache map[GceRef]*Mig

	service    *gce.Service
	cacheMutex sync.Mutex
}

// CreateGceManager constructs gceManager object.
func CreateGceManager(configReader io.Reader) (*GceManager, error) {
	// Create Google Compute Engine token.
	tokenSource := google.ComputeTokenSource("")
	if configReader != nil {
		var cfg provider_gce.Config
		if err := gcfg.ReadInto(&cfg, configReader); err != nil {
			glog.Errorf("Couldn't read config: %v", err)
			return nil, err
		}
		if cfg.Global.TokenURL == "" {
			glog.Warning("Empty tokenUrl in cloud config")
		} else {
			glog.Infof("Using TokenSource from config %#v", tokenSource)
			tokenSource = provider_gce.NewAltTokenSource(cfg.Global.TokenURL, cfg.Global.TokenBody)
		}
	} else {
		glog.Infof("Using default TokenSource %#v", tokenSource)
	}

	// Create Google Compute Engine service.
	client := oauth2.NewClient(oauth2.NoContext, tokenSource)
	gceService, err := gce.New(client)
	if err != nil {
		return nil, err
	}
	manager := &GceManager{
		migs:     make([]*migInformation, 0),
		service:  gceService,
		migCache: make(map[GceRef]*Mig),
	}
	go wait.Forever(func() {
		manager.cacheMutex.Lock()
		defer manager.cacheMutex.Unlock()
		if err := manager.regenerateCache(); err != nil {
			glog.Errorf("Error while regenerating Mig cache: %v", err)
		}
	}, time.Hour)
	return manager, nil
}

// RegisterMig registers mig in Gce Manager.
func (m *GceManager) RegisterMig(mig *Mig) {
	m.cacheMutex.Lock()
	defer m.cacheMutex.Unlock()

	m.migs = append(m.migs, &migInformation{
		config: mig,
	})
}

// GetMigSize gets MIG size.
func (m *GceManager) GetMigSize(mig *Mig) (int64, error) {
	igm, err := m.service.InstanceGroupManagers.Get(mig.Project, mig.Zone, mig.Name).Do()
	if err != nil {
		return -1, err
	}
	return igm.TargetSize, nil
}

// SetMigSize sets MIG size.
func (m *GceManager) SetMigSize(mig *Mig, size int64) error {
	op, err := m.service.InstanceGroupManagers.Resize(mig.Project, mig.Zone, mig.Name, size).Do()
	if err != nil {
		return err
	}
	if err := m.waitForOp(op, mig.Project, mig.Zone); err != nil {
		return err
	}
	return nil
}

func (m *GceManager) waitForOp(operation *gce.Operation, project string, zone string) error {
	for start := time.Now(); time.Since(start) < operationWaitTimeout; time.Sleep(operationPollInterval) {
		glog.V(4).Infof("Waiting for operation %s %s %s", project, zone, operation.Name)
		if op, err := m.service.ZoneOperations.Get(project, zone, operation.Name).Do(); err == nil {
			glog.V(4).Infof("Operation %s %s %s status: %s", project, zone, operation.Name, op.Status)
			if op.Status == "DONE" {
				return nil
			}
		} else {
			glog.Warningf("Error while getting operation %s on %s: %v", operation.Name, operation.TargetLink, err)
		}
	}
	return fmt.Errorf("Timeout while waiting for operation %s on %s to complete.", operation.Name, operation.TargetLink)
}

// DeleteInstances deletes the given instances. All instances must be controlled by the same MIG.
func (m *GceManager) DeleteInstances(instances []*GceRef) error {
	if len(instances) == 0 {
		return nil
	}
	commonMig, err := m.GetMigForInstance(instances[0])
	if err != nil {
		return err
	}
	for _, instance := range instances {
		mig, err := m.GetMigForInstance(instance)
		if err != nil {
			return err
		}
		if mig != commonMig {
			return fmt.Errorf("Connot delete instances which don't belong to the same MIG.")
		}
	}

	req := gce.InstanceGroupManagersDeleteInstancesRequest{
		Instances: []string{},
	}
	for _, instance := range instances {
		req.Instances = append(req.Instances, GenerateInstanceUrl(instance.Project, instance.Zone, instance.Name))
	}

	op, err := m.service.InstanceGroupManagers.DeleteInstances(commonMig.Project, commonMig.Zone, commonMig.Name, &req).Do()
	if err != nil {
		return err
	}
	if err := m.waitForOp(op, commonMig.Project, commonMig.Zone); err != nil {
		return err
	}
	return nil
}

// GetMigForInstance returns MigConfig of the given Instance
func (m *GceManager) GetMigForInstance(instance *GceRef) (*Mig, error) {
	m.cacheMutex.Lock()
	defer m.cacheMutex.Unlock()
	if mig, found := m.migCache[*instance]; found {
		return mig, nil
	}

	for _, mig := range m.migs {
		if mig.config.Project == instance.Project &&
			mig.config.Zone == instance.Zone &&
			strings.HasPrefix(instance.Name, mig.basename) {
			if err := m.regenerateCache(); err != nil {
				return nil, fmt.Errorf("Error while looking for MIG for instance %+v, error: %v", *instance, err)
			}
			if mig, found := m.migCache[*instance]; found {
				return mig, nil
			}
			return nil, fmt.Errorf("Instance %+v does not belong to any configured MIG", *instance)
		}
	}
	// Instance doesn't belong to any configured mig.
	return nil, nil
}

func (m *GceManager) regenerateCache() error {
	newMigCache := make(map[GceRef]*Mig)

	for _, migInfo := range m.migs {
		mig := migInfo.config
		glog.V(4).Infof("Regenerating MIG information for %s %s %s", mig.Project, mig.Zone, mig.Name)

		instanceGroupManager, err := m.service.InstanceGroupManagers.Get(mig.Project, mig.Zone, mig.Name).Do()
		if err != nil {
			return err
		}
		migInfo.basename = instanceGroupManager.BaseInstanceName

		instances, err := m.service.InstanceGroupManagers.ListManagedInstances(mig.Project, mig.Zone, mig.Name).Do()
		if err != nil {
			glog.V(4).Infof("Failed MIG info request for %s %s %s: %v", mig.Project, mig.Zone, mig.Name, err)
			return err
		}
		for _, instance := range instances.ManagedInstances {
			project, zone, name, err := ParseInstanceUrl(instance.Instance)
			if err != nil {
				return err
			}
			newMigCache[GceRef{Project: project, Zone: zone, Name: name}] = mig
		}
	}

	m.migCache = newMigCache
	return nil
}

// GetMigNodes returns mig nodes.
func (m *GceManager) GetMigNodes(mig *Mig) ([]string, error) {
	instances, err := m.service.InstanceGroupManagers.ListManagedInstances(mig.Project, mig.Zone, mig.Name).Do()
	if err != nil {
		return []string{}, err
	}
	result := make([]string, 0)
	for _, instance := range instances.ManagedInstances {
		project, zone, name, err := ParseInstanceUrl(instance.Instance)
		if err != nil {
			return []string{}, err
		}
		result = append(result, fmt.Sprintf("gce://%s/%s/%s", project, zone, name))
	}
	return result, nil
}
<|endoftext|>


<commit_before>
package main

import (
	"flag"
	"encoding/xml"
	"io/ioutil"
	"net/http"
)

type AppConfig struct {
	Quiet bool
	Listen string
	Concurrent int
	Requests uint64
	Time int
	Walks []Walk
}

type Walk struct {
	Steps []IStep
}

type IStep interface {
	Compile() error
	Run() *StepResponse
}

type TrackFile struct {
	HttpSteps []HttpStep	`xml:"httpstep"`
	Walks []TrackWalk			`xml:"walk"`
}

type TrackWalk struct {
	Name string 					`xml:"name,attr"`
	HttpSteps []HttpStep	`xml:"httpstep"`
}

type HttpStep struct {
	Name string 			`xml:"name,attr"`
	Method string 		`xml:"method,attr"`
	Url string				`xml:"url,attr"`
	Headers []Header	`xml:"header"`
	Body string				`xml:"body"`
	Request *http.Request
}

type Header struct {
	Name string				`xml:"name,attr"`
	Value string			`xml:"value,attr"`
}

func ParseArguments() []string {
	ako.AppConfig = new(AppConfig)
	flag.BoolVar(&ako.AppConfig.Quiet, "q", false, "supress messages")
	flag.StringVar(&ako.AppConfig.Listen, "l", "", "listen on this addres for instructions")
	flag.IntVar(&ako.AppConfig.Concurrent, "c", 1, "maximum concurrency")
	flag.Uint64Var(&ako.AppConfig.Requests, "n", 0, "number of requests to make")
	flag.IntVar(&ako.AppConfig.Time, "t", -1, "duration of time to send out requests")
	flag.Parse()
	return flag.Args()
}

func CompileWalkList(filepaths ... string) ([]Walk, error) {
	walks := make([]Walk, 0, 1)
	for _, path := range filepaths {
		newwalks, err := ParseInstructions(path)
		if err != nil { return nil, err }
		walks = append(walks, newwalks...)
	}
	return walks, nil
}

func ParseInstructions(path string) ([]Walk, error) {
	nw := make([]Walk, 0, 1)
	tf := new(TrackFile)
	data, ioerr := ioutil.ReadFile(path)
	if ioerr != nil { return nil, ioerr }
	xmlerr := xml.Unmarshal([]byte(data), tf)
	if xmlerr != nil { return nil, xmlerr }

	if len(tf.Walks) > 0 {
		for _, walk := range tf.Walks {
			moresteps := make([]IStep, len(walk.HttpSteps))
			for wi, wstep := range walk.HttpSteps {
				wstep.Compile()
				moresteps[wi] = &wstep
			}
			nw = append(nw, Walk{ moresteps })
		}
	} else {
		for _, step := range tf.HttpSteps {
			step.Compile()
			nw = append(nw, Walk{ []IStep{ &step } })
		}
	}

	return nw, nil
}
<commit_msg>Fixed pointer issue.<commit_after>
package main

import (
	"flag"
	"encoding/xml"
	"io/ioutil"
	"net/http"
)

type AppConfig struct {
	Quiet bool
	Listen string
	Concurrent int
	Requests uint64
	Time int
	Walks []Walk
}

type Walk struct {
	Steps []IStep
}

type IStep interface {
	Compile() error
	Run() *StepResponse
}

type TrackFile struct {
	HttpSteps []HttpStep	`xml:"httpstep"`
	Walks []TrackWalk			`xml:"walk"`
}

type TrackWalk struct {
	Name string 					`xml:"name,attr"`
	HttpSteps []HttpStep	`xml:"httpstep"`
}

type HttpStep struct {
	Name string 			`xml:"name,attr"`
	Method string 		`xml:"method,attr"`
	Url string				`xml:"url,attr"`
	Headers []Header	`xml:"header"`
	Body string				`xml:"body"`
	Request *http.Request
}

type Header struct {
	Name string				`xml:"name,attr"`
	Value string			`xml:"value,attr"`
}

func ParseArguments() []string {
	ako.AppConfig = new(AppConfig)
	flag.BoolVar(&ako.AppConfig.Quiet, "q", false, "supress messages")
	flag.StringVar(&ako.AppConfig.Listen, "l", "", "listen on this addres for instructions")
	flag.IntVar(&ako.AppConfig.Concurrent, "c", 1, "maximum concurrency")
	flag.Uint64Var(&ako.AppConfig.Requests, "n", 0, "number of requests to make")
	flag.IntVar(&ako.AppConfig.Time, "t", -1, "duration of time to send out requests")
	flag.Parse()
	return flag.Args()
}

func CompileWalkList(filepaths ... string) ([]Walk, error) {
	walks := make([]Walk, 0, 1)
	for _, path := range filepaths {
		newwalks, err := ParseInstructions(path)
		if err != nil { return nil, err }
		walks = append(walks, newwalks...)
	}
	return walks, nil
}

func ParseInstructions(path string) ([]Walk, error) {
	nw := make([]Walk, 0, 1)
	tf := new(TrackFile)
	data, ioerr := ioutil.ReadFile(path)
	if ioerr != nil { return nil, ioerr }
	xmlerr := xml.Unmarshal([]byte(data), tf)
	if xmlerr != nil { return nil, xmlerr }

	if len(tf.Walks) > 0 {
		for _, walk := range tf.Walks {
			moresteps := make([]IStep, len(walk.HttpSteps))
			for wi, wstep := range walk.HttpSteps {
				wstep.Compile()
				moresteps[wi] = &wstep
			}
			nw = append(nw, Walk{ moresteps })
		}
	} else {
		for s := range tf.HttpSteps {
			tf.HttpSteps[s].Compile()
			nw = append(nw, Walk{ []IStep{ &tf.HttpSteps[s] } })
		}
	}

	return nw, nil
}
<|endoftext|>


<commit_before>package goconfig

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"regexp"
)

// ConfigMap is shorthand for the type used as a config struct.
type ConfigMap map[string]map[string]string

var (
	configSection    = regexp.MustCompile("^\\s*\\[\\s*(\\w+)\\s*\\]\\s*$")
	quotedConfigLine = regexp.MustCompile("^\\s*(\\w+)\\s*=\\s*[\"'](.*)[\"']\\s*$")
	configLine       = regexp.MustCompile("^\\s*(\\w+)\\s*=\\s*(.*)\\s*$")
	commentLine      = regexp.MustCompile("^#.*$")
	blankLine        = regexp.MustCompile("^\\s*$")
)

var DefaultSection = "default"

// ParseFile takes the filename as a string and returns a ConfigMap.
func ParseFile(fileName string) (cfg ConfigMap, err error) {
	var file *os.File

	cfg = make(ConfigMap, 0)
	file, err = os.Open(fileName)
	if err != nil {
		return
	}
	defer file.Close()
	buf := bufio.NewReader(file)

	var (
		line           string
		longLine       bool
		currentSection string
		lineBytes      []byte
		isPrefix       bool
	)

	for {
		err = nil
		lineBytes, isPrefix, err = buf.ReadLine()
		if io.EOF == err {
			err = nil
			break
		} else if err != nil {
			break
		} else if isPrefix {
			line += string(lineBytes)

			longLine = true
			continue
		} else if longLine {
			line += string(lineBytes)
			longLine = false
		} else {
			line = string(lineBytes)
		}

		if commentLine.MatchString(line) {
			continue
		} else if blankLine.MatchString(line) {
			continue
		} else if configSection.MatchString(line) {
			section := configSection.ReplaceAllString(line,
				"$1")
			if section == "" {
				err = fmt.Errorf("invalid structure in file")
				break
			} else if !cfg.SectionInConfig(section) {
				cfg[section] = make(map[string]string, 0)
			}
			currentSection = section
		} else if configLine.MatchString(line) {
			regex := configLine
			if quotedConfigLine.MatchString(line) {
				regex = quotedConfigLine
			}
			if currentSection == "" {
				currentSection = DefaultSection
				if !cfg.SectionInConfig(currentSection) {
					cfg[currentSection] = make(map[string]string, 0)
				}
			}
			key := regex.ReplaceAllString(line, "$1")
			val := regex.ReplaceAllString(line, "$2")
			if key == "" {
				continue
			}
			cfg[currentSection][key] = val
		} else {
			err = fmt.Errorf("invalid config file")
			break
		}
	}
	return
}

// SectionInConfig determines whether a section is in the configuration.
func (c *ConfigMap) SectionInConfig(section string) bool {
	for s, _ := range *c {
		if section == s {
			return true
		}
	}
	return false
}

// ListSections returns the list of sections in the config map.
func (c *ConfigMap) ListSections() (sections []string) {
	for section, _ := range *c {
		sections = append(sections, section)
	}
	return
}

// WriteFile writes out the configuration to a file.
func (c *ConfigMap) WriteFile(filename string) (err error) {
	file, err := os.Create(filename)
	if err != nil {
		return
	}
	defer file.Close()

	for _, section := range c.ListSections() {
		sName := fmt.Sprintf("[ %s ]\n", section)
		_, err = file.Write([]byte(sName))
		if err != nil {
			return
		}

		for k, v := range (*c)[section] {
			line := fmt.Sprintf("%s = %s\n", k, v)
			_, err = file.Write([]byte(line))
			if err != nil {
				return
			}
		}
		_, err = file.Write([]byte{0x0a})
		if err != nil {
			return
		}
	}
	return
}

// AddSection creates a new section in the config map.
func (c *ConfigMap) AddSection(section string) {
	if nil != (*c)[section] {
		(*c)[section] = make(map[string]string, 0)
	}
}

// AddKeyVal adds a key value pair to a config map.
func (c *ConfigMap) AddKeyVal(section, key, val string) {
	if "" == section {
		section = DefaultSection
	}

	if nil == (*c)[section] {
		c.AddSection(section)
	}

	(*c)[section][key] = val
}

// Retrieve the value from a key map.
func (c *ConfigMap) GetValue(section, key string) (val string, present bool) {
	if c == nil {
		return
	}

	if section == "" {
		section = DefaultSection
	}

	cm := *c
	_, ok := cm[section]
	if !ok {
		return
	}

	val, present = cm[section][key]
	return
}
<commit_msg>Add a few additional functions.<commit_after>package goconfig

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"regexp"
)

// ConfigMap is shorthand for the type used as a config struct.
type ConfigMap map[string]map[string]string

var (
	configSection    = regexp.MustCompile("^\\s*\\[\\s*(\\w+)\\s*\\]\\s*$")
	quotedConfigLine = regexp.MustCompile("^\\s*(\\w+)\\s*=\\s*[\"'](.*)[\"']\\s*$")
	configLine       = regexp.MustCompile("^\\s*(\\w+)\\s*=\\s*(.*)\\s*$")
	commentLine      = regexp.MustCompile("^#.*$")
	blankLine        = regexp.MustCompile("^\\s*$")
)

var DefaultSection = "default"

// ParseFile takes the filename as a string and returns a ConfigMap.
func ParseFile(fileName string) (cfg ConfigMap, err error) {
	var file *os.File
	file, err = os.Open(fileName)
	if err != nil {
		return
	}
	defer file.Close()
	return ParseReader(file)
}

func ParseReader(r io.Reader) (cfg ConfigMap, err error) {
	cfg = make(ConfigMap, 0)
	buf := bufio.NewReader(r)

	var (
		line           string
		longLine       bool
		currentSection string
		lineBytes      []byte
		isPrefix       bool
	)

	for {
		err = nil
		lineBytes, isPrefix, err = buf.ReadLine()
		if io.EOF == err {
			err = nil
			break
		} else if err != nil {
			break
		} else if isPrefix {
			line += string(lineBytes)

			longLine = true
			continue
		} else if longLine {
			line += string(lineBytes)
			longLine = false
		} else {
			line = string(lineBytes)
		}

		if commentLine.MatchString(line) {
			continue
		} else if blankLine.MatchString(line) {
			continue
		} else if configSection.MatchString(line) {
			section := configSection.ReplaceAllString(line,
				"$1")
			if section == "" {
				err = fmt.Errorf("invalid structure in file")
				break
			} else if !cfg.SectionInConfig(section) {
				cfg[section] = make(map[string]string, 0)
			}
			currentSection = section
		} else if configLine.MatchString(line) {
			regex := configLine
			if quotedConfigLine.MatchString(line) {
				regex = quotedConfigLine
			}
			if currentSection == "" {
				currentSection = DefaultSection
				if !cfg.SectionInConfig(currentSection) {
					cfg[currentSection] = make(map[string]string, 0)
				}
			}
			key := regex.ReplaceAllString(line, "$1")
			val := regex.ReplaceAllString(line, "$2")
			if key == "" {
				continue
			}
			cfg[currentSection][key] = val
		} else {
			err = fmt.Errorf("invalid config file")
			break
		}
	}
	return
}

// SectionInConfig determines whether a section is in the configuration.
func (c *ConfigMap) SectionInConfig(section string) bool {
	for s, _ := range *c {
		if section == s {
			return true
		}
	}
	return false
}

// ListSections returns the list of sections in the config map.
func (c *ConfigMap) ListSections() (sections []string) {
	for section, _ := range *c {
		sections = append(sections, section)
	}
	return
}

// WriteFile writes out the configuration to a file.
func (c *ConfigMap) WriteFile(filename string) (err error) {
	file, err := os.Create(filename)
	if err != nil {
		return
	}
	defer file.Close()

	for _, section := range c.ListSections() {
		sName := fmt.Sprintf("[ %s ]\n", section)
		_, err = file.Write([]byte(sName))
		if err != nil {
			return
		}

		for k, v := range (*c)[section] {
			line := fmt.Sprintf("%s = %s\n", k, v)
			_, err = file.Write([]byte(line))
			if err != nil {
				return
			}
		}
		_, err = file.Write([]byte{0x0a})
		if err != nil {
			return
		}
	}
	return
}

// AddSection creates a new section in the config map.
func (c *ConfigMap) AddSection(section string) {
	if nil != (*c)[section] {
		(*c)[section] = make(map[string]string, 0)
	}
}

// AddKeyVal adds a key value pair to a config map.
func (c *ConfigMap) AddKeyVal(section, key, val string) {
	if "" == section {
		section = DefaultSection
	}

	if nil == (*c)[section] {
		c.AddSection(section)
	}

	(*c)[section][key] = val
}

// Retrieve the value from a key map.
func (c *ConfigMap) GetValue(section, key string) (val string, present bool) {
	if c == nil {
		return
	}

	if section == "" {
		section = DefaultSection
	}

	cm := *c
	_, ok := cm[section]
	if !ok {
		return
	}

	val, present = cm[section][key]
	return
}

// Return a slice of strings containing all the keys in a section.
func (c *ConfigMap) SectionKeys(section string) (keys []string, present bool) {
	if c == nil {
		return nil, false
	}

	if section == "" {
		section = DefaultSection
	}

	cm := *c
	s, ok := cm[section]
	if !ok {
		return nil, false
	}

	keys = make([]string, 0, len(s))
	for key := range s {
		keys = append(keys, key)
	}

	return keys, true
}
<|endoftext|>


<commit_before>package needle

import (
	"fmt"
	"strconv"
)

const (
	// stored unit types
	Empty byte = iota
	Minute
	Hour
	Day
	Week
	Month
	Year
)

type TTL struct {
	Count byte
	Unit  byte
}

var EMPTY_TTL = &TTL{}

// translate a readable ttl to internal ttl
// Supports format example:
// 3m: 3 minutes
// 4h: 4 hours
// 5d: 5 days
// 6w: 6 weeks
// 7M: 7 months
// 8y: 8 years
func ReadTTL(ttlString string) (*TTL, error) {
	if ttlString == "" {
		return EMPTY_TTL, nil
	}
	ttlBytes := []byte(ttlString)
	unitByte := ttlBytes[len(ttlBytes)-1]
	countBytes := ttlBytes[0 : len(ttlBytes)-1]
	if '0' <= unitByte && unitByte <= '9' {
		countBytes = ttlBytes
		unitByte = 'm'
	}
	count, err := strconv.Atoi(string(countBytes))
	unit := toStoredByte(unitByte)
	return &TTL{Count: byte(count), Unit: unit}, err
}

// read stored bytes to a ttl
func LoadTTLFromBytes(input []byte) (t *TTL) {
	if input[0] == 0 && input[1] == 0 {
		return EMPTY_TTL
	}
	return &TTL{Count: input[0], Unit: input[1]}
}

// read stored bytes to a ttl
func LoadTTLFromUint32(ttl uint32) (t *TTL) {
	input := make([]byte, 2)
	input[1] = byte(ttl)
	input[0] = byte(ttl >> 8)
	return LoadTTLFromBytes(input)
}

// save stored bytes to an output with 2 bytes
func (t *TTL) ToBytes(output []byte) {
	output[0] = t.Count
	output[1] = t.Unit
}

func (t *TTL) ToUint32() (output uint32) {
	if t == nil || t.Count == 0 {
		return 0
	}
	output = uint32(t.Count) << 8
	output += uint32(t.Unit)
	return output
}

func (t *TTL) String() string {
	if t == nil || t.Count == 0 {
		return ""
	}
	if t.Unit == Empty {
		return ""
	}
	countString := strconv.Itoa(int(t.Count))
	switch t.Unit {
	case Minute:
		return countString + "m"
	case Hour:
		return countString + "h"
	case Day:
		return countString + "d"
	case Week:
		return countString + "w"
	case Month:
		return countString + "M"
	case Year:
		return countString + "y"
	}
	return ""
}

func toStoredByte(readableUnitByte byte) byte {
	switch readableUnitByte {
	case 'm':
		return Minute
	case 'h':
		return Hour
	case 'd':
		return Day
	case 'w':
		return Week
	case 'M':
		return Month
	case 'y':
		return Year
	}
	return 0
}

func (t TTL) Minutes() uint32 {
	switch t.Unit {
	case Empty:
		return 0
	case Minute:
		return uint32(t.Count)
	case Hour:
		return uint32(t.Count) * 60
	case Day:
		return uint32(t.Count) * 60 * 24
	case Week:
		return uint32(t.Count) * 60 * 24 * 7
	case Month:
		return uint32(t.Count) * 60 * 24 * 30
	case Year:
		return uint32(t.Count) * 60 * 24 * 365
	}
	return 0
}

func SecondsToTTL(seconds int32) string {
	if seconds == 0 {
		return ""
	}
	if seconds/60 < 256 {
		return fmt.Sprintf("%dm", seconds/60)
	}
	if seconds/(3600) < 256 {
		return fmt.Sprintf("%dh", seconds/(3600))
	}
	if seconds/(3600*24) < 256 {
		return fmt.Sprintf("%dd", seconds/(3600*24))
	}
	if seconds/(3600*24*7) < 256 {
		return fmt.Sprintf("%dw", seconds/(3600*24*7))
	}
	if seconds/(3600*24*30) < 256 {
		return fmt.Sprintf("%dM", seconds/(3600*24*30))
	}
	if seconds/(3600*24*365) < 256 {
		return fmt.Sprintf("%dy", seconds/(3600*24*365))
	}
	return ""
}
<commit_msg>better unit detection<commit_after>package needle

import (
	"fmt"
	"strconv"
)

const (
	// stored unit types
	Empty byte = iota
	Minute
	Hour
	Day
	Week
	Month
	Year
)

type TTL struct {
	Count byte
	Unit  byte
}

var EMPTY_TTL = &TTL{}

// translate a readable ttl to internal ttl
// Supports format example:
// 3m: 3 minutes
// 4h: 4 hours
// 5d: 5 days
// 6w: 6 weeks
// 7M: 7 months
// 8y: 8 years
func ReadTTL(ttlString string) (*TTL, error) {
	if ttlString == "" {
		return EMPTY_TTL, nil
	}
	ttlBytes := []byte(ttlString)
	unitByte := ttlBytes[len(ttlBytes)-1]
	countBytes := ttlBytes[0 : len(ttlBytes)-1]
	if '0' <= unitByte && unitByte <= '9' {
		countBytes = ttlBytes
		unitByte = 'm'
	}
	count, err := strconv.Atoi(string(countBytes))
	unit := toStoredByte(unitByte)
	return &TTL{Count: byte(count), Unit: unit}, err
}

// read stored bytes to a ttl
func LoadTTLFromBytes(input []byte) (t *TTL) {
	if input[0] == 0 && input[1] == 0 {
		return EMPTY_TTL
	}
	return &TTL{Count: input[0], Unit: input[1]}
}

// read stored bytes to a ttl
func LoadTTLFromUint32(ttl uint32) (t *TTL) {
	input := make([]byte, 2)
	input[1] = byte(ttl)
	input[0] = byte(ttl >> 8)
	return LoadTTLFromBytes(input)
}

// save stored bytes to an output with 2 bytes
func (t *TTL) ToBytes(output []byte) {
	output[0] = t.Count
	output[1] = t.Unit
}

func (t *TTL) ToUint32() (output uint32) {
	if t == nil || t.Count == 0 {
		return 0
	}
	output = uint32(t.Count) << 8
	output += uint32(t.Unit)
	return output
}

func (t *TTL) String() string {
	if t == nil || t.Count == 0 {
		return ""
	}
	if t.Unit == Empty {
		return ""
	}
	countString := strconv.Itoa(int(t.Count))
	switch t.Unit {
	case Minute:
		return countString + "m"
	case Hour:
		return countString + "h"
	case Day:
		return countString + "d"
	case Week:
		return countString + "w"
	case Month:
		return countString + "M"
	case Year:
		return countString + "y"
	}
	return ""
}

func toStoredByte(readableUnitByte byte) byte {
	switch readableUnitByte {
	case 'm':
		return Minute
	case 'h':
		return Hour
	case 'd':
		return Day
	case 'w':
		return Week
	case 'M':
		return Month
	case 'y':
		return Year
	}
	return 0
}

func (t TTL) Minutes() uint32 {
	switch t.Unit {
	case Empty:
		return 0
	case Minute:
		return uint32(t.Count)
	case Hour:
		return uint32(t.Count) * 60
	case Day:
		return uint32(t.Count) * 60 * 24
	case Week:
		return uint32(t.Count) * 60 * 24 * 7
	case Month:
		return uint32(t.Count) * 60 * 24 * 30
	case Year:
		return uint32(t.Count) * 60 * 24 * 365
	}
	return 0
}

func SecondsToTTL(seconds int32) string {
	if seconds == 0 {
		return ""
	}
	if seconds%(3600*24*365) == 0 && seconds/(3600*24*365) < 256 {
		return fmt.Sprintf("%dy", seconds/(3600*24*365))
	}
	if seconds%(3600*24*30) == 0 && seconds/(3600*24*30) < 256 {
		return fmt.Sprintf("%dM", seconds/(3600*24*30))
	}
	if seconds%(3600*24*7) == 0 && seconds/(3600*24*7) < 256 {
		return fmt.Sprintf("%dw", seconds/(3600*24*7))
	}
	if seconds%(3600*24) == 0 && seconds/(3600*24) < 256 {
		return fmt.Sprintf("%dd", seconds/(3600*24))
	}
	if seconds%(3600) == 0 && seconds/(3600) < 256 {
		return fmt.Sprintf("%dh", seconds/(3600))
	}
	if seconds/60 < 256 {
		return fmt.Sprintf("%dm", seconds/60)
	}
	if seconds/(3600) < 256 {
		return fmt.Sprintf("%dh", seconds/(3600))
	}
	if seconds/(3600*24) < 256 {
		return fmt.Sprintf("%dd", seconds/(3600*24))
	}
	if seconds/(3600*24*7) < 256 {
		return fmt.Sprintf("%dw", seconds/(3600*24*7))
	}
	if seconds/(3600*24*30) < 256 {
		return fmt.Sprintf("%dM", seconds/(3600*24*30))
	}
	if seconds/(3600*24*365) < 256 {
		return fmt.Sprintf("%dy", seconds/(3600*24*365))
	}
	return ""
}
<|endoftext|>


<commit_before>package helpers

import (
	"time"

	"regexp"

	"strings"

	"fmt"

	"github.com/Seklfreak/Robyul2/cache"
	"github.com/Seklfreak/Robyul2/models"
	"github.com/bwmarrin/discordgo"
	"github.com/globalsign/mgo/bson"
	redisCache "github.com/go-redis/cache"
	"github.com/pkg/errors"
)

var isValidVanityName = regexp.MustCompile(`^[a-zA-Z0-9]+$`).MatchString

func UpdateOrInsertVanityUrl(vanityName, guildID, channelID, userID string) (err error) {
	if !isValidVanityName(vanityName) {
		return errors.New("invalid vanity name")
	}

	vanityEntryByName, _ := GetVanityUrlByVanityName(vanityName)
	if vanityEntryByName.VanityName != "" && vanityEntryByName.GuildID != guildID {
		return errors.New("vanity name already in use")
	}

	vanityEntryByGuildID, _ := GetVanityUrlByGuildID(guildID)
	if vanityEntryByGuildID.VanityName != "" {
		// update vanity invite
		oldVanityInvite := vanityEntryByGuildID

		vanityEntryByGuildID.VanityName = strings.ToLower(vanityName)
		vanityEntryByGuildID.VanityNamePretty = vanityName
		vanityEntryByGuildID.ChannelID = channelID
		vanityEntryByGuildID.SetAt = time.Now()
		vanityEntryByGuildID.SetByUserID = userID

		err = ResetCachedDiscordInviteByVanityInvite(vanityEntryByGuildID)
		RelaxLog(err)

		cache.GetLogger().WithField("module", "helpers/vanityinvite").
			WithField("guildID", vanityEntryByGuildID.GuildID).
			Info("updated the vanity url, new name:", vanityEntryByGuildID.VanityNamePretty, "new channel:", vanityEntryByGuildID.ChannelID)

		err = UpdateVanityUrl(vanityEntryByGuildID)
		if err != nil {
			return err
		}

		go func() {
			logChannelID, _ := GetBotConfigString(models.VanityInviteLogChannelKey)
			if logChannelID != "" {
				err = logVanityInviteChange(logChannelID, vanityEntryByGuildID, oldVanityInvite)
				RelaxLog(err)
			}
		}()

		_, err = EventlogLog(time.Now(), guildID, guildID,
			models.EventlogTargetTypeGuild, userID,
			models.EventlogTypeRobyulVanityInviteUpdate, "",
			[]models.ElasticEventlogChange{
				{
					Key:      "vanityinvite_name",
					OldValue: oldVanityInvite.VanityName,
					NewValue: vanityEntryByGuildID.VanityName,
				},
				{
					Key:      "vanityinvite_channelid",
					OldValue: oldVanityInvite.ChannelID,
					NewValue: vanityEntryByGuildID.ChannelID,
					Type:     models.EventlogTargetTypeChannel,
				},
			},
			nil, false)
		RelaxLog(err)

		return nil
	}

	// create vanity invite
	newVanityInvite := models.VanityInviteEntry{
		GuildID:          guildID,
		VanityName:       strings.ToLower(vanityName),
		VanityNamePretty: vanityName,
		SetByUserID:      userID,
		SetAt:            time.Now(),
		ChannelID:        channelID,
	}

	cache.GetLogger().WithField("module", "helpers/vanityinvite").
		WithField("guildID", vanityEntryByGuildID.GuildID).
		Info("created the vanity url, name:", vanityEntryByGuildID.VanityNamePretty, "channel:", vanityEntryByGuildID.ChannelID)

	_, err = MDbInsert(
		models.VanityInvitesTable,
		newVanityInvite,
	)
	if err != nil {
		return err
	}

	go func() {
		logChannelID, _ := GetBotConfigString(models.VanityInviteLogChannelKey)
		if logChannelID != "" {
			err = logVanityInviteCreation(logChannelID, newVanityInvite)
			RelaxLog(err)
		}
	}()

	_, err = EventlogLog(time.Now(), guildID, guildID,
		models.EventlogTargetTypeGuild, userID,
		models.EventlogTypeRobyulVanityInviteCreate, "",
		nil,
		[]models.ElasticEventlogOption{
			{
				Key:   "vanityinvite_name",
				Value: vanityName,
			},
			{
				Key:   "vanityinvite_channelid",
				Value: channelID,
				Type:  models.EventlogTargetTypeChannel,
			},
		}, false)
	RelaxLog(err)

	err = ResetCachedDiscordInviteByVanityInvite(newVanityInvite)
	RelaxLog(err)

	return nil
}

func logVanityInviteCreation(channelID string, vanityInvite models.VanityInviteEntry) (err error) {
	author, err := GetUser(vanityInvite.SetByUserID)
	if err != nil {
		return err
	}

	newChannel, err := GetChannel(vanityInvite.ChannelID)
	if err != nil {
		newChannel = new(discordgo.Channel)
		newChannel.Name = "N/A"
	}

	newGuild, err := GetGuild(vanityInvite.GuildID)
	if err != nil {
		newGuild = new(discordgo.Guild)
		newGuild.Name = "N/A"
	}

	embed := &discordgo.MessageEmbed{
		Author: &discordgo.MessageEmbedAuthor{
			Name: fmt.Sprintf("@%s#%s created a Custom Invite 🆕:", author.Username, author.Discriminator),
		},
		Color: 0x0FADED,
		Footer: &discordgo.MessageEmbedFooter{
			Text: fmt.Sprintf("Vanity Invite ID: %s | Set at %s",
				vanityInvite.ID, vanityInvite.SetAt.Format(time.ANSIC)),
		},
		Fields: []*discordgo.MessageEmbedField{
			{
				Name:   "Name",
				Value:  fmt.Sprintf("`%s`", vanityInvite.VanityNamePretty),
				Inline: false,
			},
			{
				Name: "Channel",
				Value: fmt.Sprintf("`#%s`\n`#%s`",
					newChannel.Name, vanityInvite.ChannelID),
				Inline: false,
			},
			{
				Name: "Server",
				Value: fmt.Sprintf("`%s`\n`#%s`",
					newGuild.Name, vanityInvite.GuildID),
				Inline: false,
			},
		},
	}

	if author.Avatar != "" {
		embed.Author.IconURL = author.AvatarURL("")
	}

	_, err = SendEmbed(channelID, embed)
	return err
}

func logVanityInviteDeletion(channelID string, vanityInvite models.VanityInviteEntry) (err error) {
	author, err := GetUser(vanityInvite.SetByUserID)
	if err != nil {
		return err
	}

	newChannel, err := GetChannel(vanityInvite.ChannelID)
	if err != nil {
		newChannel = new(discordgo.Channel)
		newChannel.Name = "N/A"
	}

	newGuild, err := GetGuild(vanityInvite.GuildID)
	if err != nil {
		newGuild = new(discordgo.Guild)
		newGuild.Name = "N/A"
	}

	embed := &discordgo.MessageEmbed{
		Author: &discordgo.MessageEmbedAuthor{
			Name: fmt.Sprintf("@%s#%s deleted a Custom Invite 🚮:", author.Username, author.Discriminator),
		},
		Color: 0x0FADED,
		Footer: &discordgo.MessageEmbedFooter{
			Text: fmt.Sprintf("Vanity Invite ID: %s | Set at %s",
				vanityInvite.ID, vanityInvite.SetAt.Format(time.ANSIC)),
		},
		Fields: []*discordgo.MessageEmbedField{
			{
				Name:   "Name",
				Value:  fmt.Sprintf("`%s`", vanityInvite.VanityNamePretty),
				Inline: false,
			},
			{
				Name: "Channel",
				Value: fmt.Sprintf("`#%s`\n`#%s`",
					newChannel.Name, vanityInvite.ChannelID),
				Inline: false,
			},
			{
				Name: "Server",
				Value: fmt.Sprintf("`%s`\n`#%s`",
					newGuild.Name, vanityInvite.GuildID),
				Inline: false,
			},
		},
	}

	if author.Avatar != "" {
		embed.Author.IconURL = author.AvatarURL("")
	}

	_, err = SendEmbed(channelID, embed)
	return err
}

func logVanityInviteChange(channelID string, vanityInvite, oldVanityInvite models.VanityInviteEntry) (err error) {
	author, err := GetUser(vanityInvite.SetByUserID)
	if err != nil {
		return err
	}

	newGuild, err := GetGuild(vanityInvite.GuildID)
	if err != nil {
		newGuild = new(discordgo.Guild)
		newGuild.Name = "N/A"
	}

	embed := &discordgo.MessageEmbed{
		Author: &discordgo.MessageEmbedAuthor{
			Name: fmt.Sprintf("@%s#%s changed a Custom Invite 🔄:", author.Username, author.Discriminator),
		},
		Color: 0x0FADED,
		Footer: &discordgo.MessageEmbedFooter{
			Text: fmt.Sprintf("Vanity Invite ID: %s | Set at %s",
				vanityInvite.ID, vanityInvite.SetAt.Format(time.ANSIC)),
		},
		Fields: []*discordgo.MessageEmbedField{},
	}

	if vanityInvite.VanityNamePretty != oldVanityInvite.VanityNamePretty {
		embed.Fields = append(embed.Fields, &discordgo.MessageEmbedField{
			Name:   "Name",
			Value:  fmt.Sprintf("`%s` ➡ `%s`", oldVanityInvite.VanityNamePretty, vanityInvite.VanityNamePretty),
			Inline: false,
		})
	}
	if vanityInvite.ChannelID != oldVanityInvite.ChannelID {
		newChannel, err := GetChannel(vanityInvite.ChannelID)
		if err != nil {
			newChannel = new(discordgo.Channel)
			newChannel.Name = "N/A"
		}
		oldChannel, err := GetChannel(oldVanityInvite.ChannelID)
		if err != nil {
			oldChannel = new(discordgo.Channel)
			oldChannel.Name = "N/A"
		}

		embed.Fields = append(embed.Fields, &discordgo.MessageEmbedField{
			Name: "Channel",
			Value: fmt.Sprintf("`#%s` ➡ `#%s`\n`#%s` ➡ `#%s`",
				oldChannel.Name, newChannel.Name,
				oldVanityInvite.ChannelID, vanityInvite.ChannelID),
			Inline: false,
		})
	}

	embed.Fields = append(embed.Fields, &discordgo.MessageEmbedField{
		Name: "Server",
		Value: fmt.Sprintf("`%s`\n`#%s`",
			newGuild.Name, vanityInvite.GuildID),
		Inline: true,
	})

	if author.Avatar != "" {
		embed.Author.IconURL = author.AvatarURL("")
	}

	_, err = SendEmbed(channelID, embed)
	return err
}

func GetVanityUrlByGuildID(guildID string) (entryBucket models.VanityInviteEntry, err error) {
	err = MdbOneWithoutLogging(
		MdbCollection(models.VanityInvitesTable).Find(bson.M{"guildid": guildID}),
		&entryBucket,
	)
	return entryBucket, err
}

func GetVanityUrlByVanityName(vanityName string) (entryBucket models.VanityInviteEntry, err error) {
	err = MdbOneWithoutLogging(
		MdbCollection(models.VanityInvitesTable).Find(bson.M{"vanityname": strings.ToLower(vanityName)}),
		&entryBucket,
	)
	return entryBucket, err
}

func RemoveVanityUrl(vanityInviteEntry models.VanityInviteEntry) error {
	if vanityInviteEntry.ID.Valid() {
		cache.GetLogger().WithField("module", "helpers/vanityinvite").
			WithField("guildID", vanityInviteEntry.GuildID).
			Info("removed the vanity url", vanityInviteEntry.VanityNamePretty)

		err := MDbDelete(models.VanityInvitesTable, vanityInviteEntry.ID)
		if err != nil {
			return err
		}

		go func() {
			logChannelID, _ := GetBotConfigString(models.VanityInviteLogChannelKey)
			if logChannelID != "" {
				err = logVanityInviteDeletion(logChannelID, vanityInviteEntry)
				RelaxLog(err)
			}
		}()

		return nil
	}
	return errors.New("empty vanityName submitted")
}

func UpdateVanityUrl(vanityInviteEntry models.VanityInviteEntry) error {
	if vanityInviteEntry.ID.Valid() {
		err := MDbUpdate(models.VanityInvitesTable, vanityInviteEntry.ID, vanityInviteEntry)
		return err
	}
	return errors.New("empty vanityName submitted")
}

func ResetCachedDiscordInviteByVanityInvite(vanityInviteEntry models.VanityInviteEntry) (err error) {
	cacheCodec := cache.GetRedisCacheCodec()
	key := fmt.Sprintf(models.VanityInvitesInviteRedisKey, vanityInviteEntry.GuildID)

	err = cacheCodec.Delete(key)
	if err != nil {
		if !strings.Contains(err.Error(), "cache: key is missing") {
			RelaxLog(err)
		}
	}

	_, err = GetDiscordInviteByVanityInvite(vanityInviteEntry)
	return err
}

func GetDiscordInviteByVanityInvite(vanityInviteEntry models.VanityInviteEntry) (code string, err error) {
	cacheCodec := cache.GetRedisCacheCodec()
	key := fmt.Sprintf(models.VanityInvitesInviteRedisKey, vanityInviteEntry.GuildID)

	var vanityInviteRedis models.VanityInviteRedisEntry
	if err = cacheCodec.Get(key, &vanityInviteRedis); err == nil {
		if time.Now().Before(vanityInviteRedis.ExpiresAt) {

			//fmt.Println(vanityInviteRedis.ExpiresAt)
			return vanityInviteRedis.InviteCode, nil
		}
	}

	cache.GetLogger().WithField("module", "vanityinvite").Infof(
		"created a new invite for Guild #%s to Channel #%s", vanityInviteEntry.GuildID, vanityInviteEntry.ChannelID)

	invite, err := cache.GetSession().ChannelInviteCreate(vanityInviteEntry.ChannelID, discordgo.Invite{
		MaxAge: 60 * 60 * 24, // 1 day
		//Unique: true,
	})
	if err != nil {
		return "", err
	}

	createdAtTime, err := invite.CreatedAt.Parse()
	if err != nil {
		createdAtTime = time.Now()
	}

	vanityInviteRedis = models.VanityInviteRedisEntry{
		InviteCode: invite.Code,
		ExpiresAt:  createdAtTime.Add(time.Duration(invite.MaxAge) * time.Second).Add(-3 * time.Hour),
	}

	err = cacheCodec.Set(&redisCache.Item{
		Key:        key,
		Object:     vanityInviteRedis,
		Expiration: createdAtTime.Add(time.Duration(invite.MaxAge) * time.Second).Add(-3 * time.Hour).Sub(time.Now()),
	})
	RelaxLog(err)

	return invite.Code, nil
}
<commit_msg>[vanityinvite] improve caching 🔗<commit_after>package helpers

import (
	"encoding/json"
	"time"

	"regexp"

	"strings"

	"fmt"

	"github.com/Seklfreak/Robyul2/cache"
	"github.com/Seklfreak/Robyul2/models"
	"github.com/bwmarrin/discordgo"
	"github.com/globalsign/mgo/bson"
	"github.com/pkg/errors"
)

var isValidVanityName = regexp.MustCompile(`^[a-zA-Z0-9]+$`).MatchString

func UpdateOrInsertVanityUrl(vanityName, guildID, channelID, userID string) (err error) {
	if !isValidVanityName(vanityName) {
		return errors.New("invalid vanity name")
	}

	vanityEntryByName, _ := GetVanityUrlByVanityName(vanityName)
	if vanityEntryByName.VanityName != "" && vanityEntryByName.GuildID != guildID {
		return errors.New("vanity name already in use")
	}

	vanityEntryByGuildID, _ := GetVanityUrlByGuildID(guildID)
	if vanityEntryByGuildID.VanityName != "" {
		// update vanity invite
		oldVanityInvite := vanityEntryByGuildID

		vanityEntryByGuildID.VanityName = strings.ToLower(vanityName)
		vanityEntryByGuildID.VanityNamePretty = vanityName
		vanityEntryByGuildID.ChannelID = channelID
		vanityEntryByGuildID.SetAt = time.Now()
		vanityEntryByGuildID.SetByUserID = userID

		err = ResetCachedDiscordInviteByVanityInvite(vanityEntryByGuildID)
		RelaxLog(err)

		cache.GetLogger().WithField("module", "helpers/vanityinvite").
			WithField("guildID", vanityEntryByGuildID.GuildID).
			Info("updated the vanity url, new name:", vanityEntryByGuildID.VanityNamePretty, "new channel:", vanityEntryByGuildID.ChannelID)

		err = UpdateVanityUrl(vanityEntryByGuildID)
		if err != nil {
			return err
		}

		go func() {
			logChannelID, _ := GetBotConfigString(models.VanityInviteLogChannelKey)
			if logChannelID != "" {
				err = logVanityInviteChange(logChannelID, vanityEntryByGuildID, oldVanityInvite)
				RelaxLog(err)
			}
		}()

		_, err = EventlogLog(time.Now(), guildID, guildID,
			models.EventlogTargetTypeGuild, userID,
			models.EventlogTypeRobyulVanityInviteUpdate, "",
			[]models.ElasticEventlogChange{
				{
					Key:      "vanityinvite_name",
					OldValue: oldVanityInvite.VanityName,
					NewValue: vanityEntryByGuildID.VanityName,
				},
				{
					Key:      "vanityinvite_channelid",
					OldValue: oldVanityInvite.ChannelID,
					NewValue: vanityEntryByGuildID.ChannelID,
					Type:     models.EventlogTargetTypeChannel,
				},
			},
			nil, false)
		RelaxLog(err)

		return nil
	}

	// create vanity invite
	newVanityInvite := models.VanityInviteEntry{
		GuildID:          guildID,
		VanityName:       strings.ToLower(vanityName),
		VanityNamePretty: vanityName,
		SetByUserID:      userID,
		SetAt:            time.Now(),
		ChannelID:        channelID,
	}

	cache.GetLogger().WithField("module", "helpers/vanityinvite").
		WithField("guildID", vanityEntryByGuildID.GuildID).
		Info("created the vanity url, name:", vanityEntryByGuildID.VanityNamePretty, "channel:", vanityEntryByGuildID.ChannelID)

	_, err = MDbInsert(
		models.VanityInvitesTable,
		newVanityInvite,
	)
	if err != nil {
		return err
	}

	go func() {
		logChannelID, _ := GetBotConfigString(models.VanityInviteLogChannelKey)
		if logChannelID != "" {
			err = logVanityInviteCreation(logChannelID, newVanityInvite)
			RelaxLog(err)
		}
	}()

	_, err = EventlogLog(time.Now(), guildID, guildID,
		models.EventlogTargetTypeGuild, userID,
		models.EventlogTypeRobyulVanityInviteCreate, "",
		nil,
		[]models.ElasticEventlogOption{
			{
				Key:   "vanityinvite_name",
				Value: vanityName,
			},
			{
				Key:   "vanityinvite_channelid",
				Value: channelID,
				Type:  models.EventlogTargetTypeChannel,
			},
		}, false)
	RelaxLog(err)

	err = ResetCachedDiscordInviteByVanityInvite(newVanityInvite)
	RelaxLog(err)

	return nil
}

func logVanityInviteCreation(channelID string, vanityInvite models.VanityInviteEntry) (err error) {
	author, err := GetUser(vanityInvite.SetByUserID)
	if err != nil {
		return err
	}

	newChannel, err := GetChannel(vanityInvite.ChannelID)
	if err != nil {
		newChannel = new(discordgo.Channel)
		newChannel.Name = "N/A"
	}

	newGuild, err := GetGuild(vanityInvite.GuildID)
	if err != nil {
		newGuild = new(discordgo.Guild)
		newGuild.Name = "N/A"
	}

	embed := &discordgo.MessageEmbed{
		Author: &discordgo.MessageEmbedAuthor{
			Name: fmt.Sprintf("@%s#%s created a Custom Invite 🆕:", author.Username, author.Discriminator),
		},
		Color: 0x0FADED,
		Footer: &discordgo.MessageEmbedFooter{
			Text: fmt.Sprintf("Vanity Invite ID: %s | Set at %s",
				vanityInvite.ID, vanityInvite.SetAt.Format(time.ANSIC)),
		},
		Fields: []*discordgo.MessageEmbedField{
			{
				Name:   "Name",
				Value:  fmt.Sprintf("`%s`", vanityInvite.VanityNamePretty),
				Inline: false,
			},
			{
				Name: "Channel",
				Value: fmt.Sprintf("`#%s`\n`#%s`",
					newChannel.Name, vanityInvite.ChannelID),
				Inline: false,
			},
			{
				Name: "Server",
				Value: fmt.Sprintf("`%s`\n`#%s`",
					newGuild.Name, vanityInvite.GuildID),
				Inline: false,
			},
		},
	}

	if author.Avatar != "" {
		embed.Author.IconURL = author.AvatarURL("")
	}

	_, err = SendEmbed(channelID, embed)
	return err
}

func logVanityInviteDeletion(channelID string, vanityInvite models.VanityInviteEntry) (err error) {
	author, err := GetUser(vanityInvite.SetByUserID)
	if err != nil {
		return err
	}

	newChannel, err := GetChannel(vanityInvite.ChannelID)
	if err != nil {
		newChannel = new(discordgo.Channel)
		newChannel.Name = "N/A"
	}

	newGuild, err := GetGuild(vanityInvite.GuildID)
	if err != nil {
		newGuild = new(discordgo.Guild)
		newGuild.Name = "N/A"
	}

	embed := &discordgo.MessageEmbed{
		Author: &discordgo.MessageEmbedAuthor{
			Name: fmt.Sprintf("@%s#%s deleted a Custom Invite 🚮:", author.Username, author.Discriminator),
		},
		Color: 0x0FADED,
		Footer: &discordgo.MessageEmbedFooter{
			Text: fmt.Sprintf("Vanity Invite ID: %s | Set at %s",
				vanityInvite.ID, vanityInvite.SetAt.Format(time.ANSIC)),
		},
		Fields: []*discordgo.MessageEmbedField{
			{
				Name:   "Name",
				Value:  fmt.Sprintf("`%s`", vanityInvite.VanityNamePretty),
				Inline: false,
			},
			{
				Name: "Channel",
				Value: fmt.Sprintf("`#%s`\n`#%s`",
					newChannel.Name, vanityInvite.ChannelID),
				Inline: false,
			},
			{
				Name: "Server",
				Value: fmt.Sprintf("`%s`\n`#%s`",
					newGuild.Name, vanityInvite.GuildID),
				Inline: false,
			},
		},
	}

	if author.Avatar != "" {
		embed.Author.IconURL = author.AvatarURL("")
	}

	_, err = SendEmbed(channelID, embed)
	return err
}

func logVanityInviteChange(channelID string, vanityInvite, oldVanityInvite models.VanityInviteEntry) (err error) {
	author, err := GetUser(vanityInvite.SetByUserID)
	if err != nil {
		return err
	}

	newGuild, err := GetGuild(vanityInvite.GuildID)
	if err != nil {
		newGuild = new(discordgo.Guild)
		newGuild.Name = "N/A"
	}

	embed := &discordgo.MessageEmbed{
		Author: &discordgo.MessageEmbedAuthor{
			Name: fmt.Sprintf("@%s#%s changed a Custom Invite 🔄:", author.Username, author.Discriminator),
		},
		Color: 0x0FADED,
		Footer: &discordgo.MessageEmbedFooter{
			Text: fmt.Sprintf("Vanity Invite ID: %s | Set at %s",
				vanityInvite.ID, vanityInvite.SetAt.Format(time.ANSIC)),
		},
		Fields: []*discordgo.MessageEmbedField{},
	}

	if vanityInvite.VanityNamePretty != oldVanityInvite.VanityNamePretty {
		embed.Fields = append(embed.Fields, &discordgo.MessageEmbedField{
			Name:   "Name",
			Value:  fmt.Sprintf("`%s` ➡ `%s`", oldVanityInvite.VanityNamePretty, vanityInvite.VanityNamePretty),
			Inline: false,
		})
	}
	if vanityInvite.ChannelID != oldVanityInvite.ChannelID {
		newChannel, err := GetChannel(vanityInvite.ChannelID)
		if err != nil {
			newChannel = new(discordgo.Channel)
			newChannel.Name = "N/A"
		}
		oldChannel, err := GetChannel(oldVanityInvite.ChannelID)
		if err != nil {
			oldChannel = new(discordgo.Channel)
			oldChannel.Name = "N/A"
		}

		embed.Fields = append(embed.Fields, &discordgo.MessageEmbedField{
			Name: "Channel",
			Value: fmt.Sprintf("`#%s` ➡ `#%s`\n`#%s` ➡ `#%s`",
				oldChannel.Name, newChannel.Name,
				oldVanityInvite.ChannelID, vanityInvite.ChannelID),
			Inline: false,
		})
	}

	embed.Fields = append(embed.Fields, &discordgo.MessageEmbedField{
		Name: "Server",
		Value: fmt.Sprintf("`%s`\n`#%s`",
			newGuild.Name, vanityInvite.GuildID),
		Inline: true,
	})

	if author.Avatar != "" {
		embed.Author.IconURL = author.AvatarURL("")
	}

	_, err = SendEmbed(channelID, embed)
	return err
}

func GetVanityUrlByGuildID(guildID string) (entryBucket models.VanityInviteEntry, err error) {
	err = MdbOneWithoutLogging(
		MdbCollection(models.VanityInvitesTable).Find(bson.M{"guildid": guildID}),
		&entryBucket,
	)
	return entryBucket, err
}

func GetVanityUrlByVanityName(vanityName string) (entryBucket models.VanityInviteEntry, err error) {
	err = MdbOneWithoutLogging(
		MdbCollection(models.VanityInvitesTable).Find(bson.M{"vanityname": strings.ToLower(vanityName)}),
		&entryBucket,
	)
	return entryBucket, err
}

func RemoveVanityUrl(vanityInviteEntry models.VanityInviteEntry) error {
	if vanityInviteEntry.ID.Valid() {
		cache.GetLogger().WithField("module", "helpers/vanityinvite").
			WithField("guildID", vanityInviteEntry.GuildID).
			Info("removed the vanity url", vanityInviteEntry.VanityNamePretty)

		err := MDbDelete(models.VanityInvitesTable, vanityInviteEntry.ID)
		if err != nil {
			return err
		}

		go func() {
			logChannelID, _ := GetBotConfigString(models.VanityInviteLogChannelKey)
			if logChannelID != "" {
				err = logVanityInviteDeletion(logChannelID, vanityInviteEntry)
				RelaxLog(err)
			}
		}()

		return nil
	}
	return errors.New("empty vanityName submitted")
}

func UpdateVanityUrl(vanityInviteEntry models.VanityInviteEntry) error {
	if vanityInviteEntry.ID.Valid() {
		err := MDbUpdate(models.VanityInvitesTable, vanityInviteEntry.ID, vanityInviteEntry)
		return err
	}
	return errors.New("empty vanityName submitted")
}

func ResetCachedDiscordInviteByVanityInvite(vanityInviteEntry models.VanityInviteEntry) (err error) {
	cacheCodec := cache.GetRedisCacheCodec()
	key := fmt.Sprintf(models.VanityInvitesInviteRedisKey, vanityInviteEntry.GuildID)

	err = cacheCodec.Delete(key)
	if err != nil {
		if !strings.Contains(err.Error(), "cache: key is missing") {
			RelaxLog(err)
		}
	}

	_, err = GetDiscordInviteByVanityInvite(vanityInviteEntry)
	return err
}

func GetDiscordInviteByVanityInvite(vanityInviteEntry models.VanityInviteEntry) (code string, err error) {
	redisClient := cache.GetRedisClient()
	key := fmt.Sprintf(models.VanityInvitesInviteRedisKey, vanityInviteEntry.GuildID)

	var vanityInviteRedis models.VanityInviteRedisEntry

	cacheResult, err := redisClient.Get(key).Bytes()
	if err == nil {
		err = json.Unmarshal(cacheResult, &vanityInviteRedis)
		if err == nil && time.Now().Before(vanityInviteRedis.ExpiresAt) {

			fmt.Printf("got item from cache: %#v, %s\n", vanityInviteRedis, vanityInviteRedis.ExpiresAt.String()) // TODO
			return vanityInviteRedis.InviteCode, nil
		}
	}

	cache.GetLogger().WithField("module", "vanityinvite").Infof(
		"created a new invite for Guild #%s to Channel #%s", vanityInviteEntry.GuildID, vanityInviteEntry.ChannelID)

	invite, err := cache.GetSession().ChannelInviteCreate(vanityInviteEntry.ChannelID, discordgo.Invite{
		MaxAge: 60 * 60 * 24, // 1 day
		// Unique: true,
	})
	if err != nil {
		return "", err
	}

	createdAtTime, err := invite.CreatedAt.Parse()
	if err != nil {
		createdAtTime = time.Now()
	}

	expireAt := createdAtTime.Add(time.Duration(invite.MaxAge) * time.Second).Add(-3 * time.Hour)

	vanityInviteRedis = models.VanityInviteRedisEntry{
		InviteCode: invite.Code,
		ExpiresAt:  expireAt,
	}

	fmt.Printf("caching item: %#v, %s, %s\n", vanityInviteRedis, expireAt.String(), expireAt.Sub(time.Now()).String()) // TODO

	cacheItem, err := json.Marshal(vanityInviteRedis)
	if err == nil {
		err = redisClient.Set(key, cacheItem, expireAt.Sub(time.Now())).Err()
		RelaxLog(err)
	}

	return invite.Code, nil
}
<|endoftext|>


<commit_before><commit_msg>Ident RFC citation<commit_after><|endoftext|>