<commit_before>package artifice

import (
	"context"
	"fmt"
	"log"
	"time"

	"github.com/antihax/evedata/internal/redisqueue"
	"github.com/antihax/evedata/internal/sqlhelper"
)

func init() {
	registerTrigger("characterTransactions", characterTransactions, time.NewTicker(time.Second*3600))
	registerTrigger("characterAssets", characterAssets, time.NewTicker(time.Second*3600))
	registerTrigger("characterOrders", characterOrders, time.NewTicker(time.Second*1200))
	registerTrigger("characterStructures", characterStructures, time.NewTicker(time.Second*3600))
	registerTrigger("characterStructureMarket", characterStructureMarket, time.NewTicker(time.Second*300))

	registerTrigger("characterNotifications", characterNotifications, time.NewTicker(time.Second*600))
	registerTrigger("characterContactSync", characterContactSync, time.NewTicker(time.Second*360))
	registerTrigger("characterAuthOwners", characterAuthOwners, time.NewTicker(time.Second*3600))
	registerTrigger("crestCharacters", crestCharacters, time.NewTicker(time.Second*3600))

	registerTrigger("allianceContacts", allianceContacts, time.NewTicker(time.Second*360))
	registerTrigger("corporationContacts", corporationContacts, time.NewTicker(time.Second*360))
}

func characterTransactions(s *Artifice) error {
	if pairs, err := s.GetCharactersForScope("read_character_wallet"); err != nil {
		return err
	} else {
		work := []redisqueue.Work{}
		for _, p := range pairs {
			work = append(work, redisqueue.Work{Operation: "characterWalletTransactions", Parameter: []int32{p.CharacterID, p.TokenCharacterID}})
			work = append(work, redisqueue.Work{Operation: "characterWalletJournal", Parameter: []int32{p.CharacterID, p.TokenCharacterID}})
		}
		return s.QueueWork(work, redisqueue.Priority_Normal)
	}
}

func characterAssets(s *Artifice) error {
	work := []redisqueue.Work{}
	if pairs, err := s.GetCharactersForScope("read_assets"); err != nil {
		return err
	} else {
		for _, p := range pairs {
			work = append(work, redisqueue.Work{Operation: "characterAssets", Parameter: []int32{p.CharacterID, p.TokenCharacterID}})
		}
	}

	return s.QueueWork(work, redisqueue.Priority_High)
}

func characterStructures(s *Artifice) error {
	structures, err := s.db.Query(
		`SELECT T.characterID, T.tokenCharacterID, S.locationID
			FROM evedata.crestTokens T
			INNER JOIN (
					SELECT DISTINCT characterID, locationID FROM evedata.orders WHERE locationID > 70000000
					UNION DISTINCT
					SELECT DISTINCT characterID, locationID FROM evedata.assets WHERE locationID > 70000000
						AND locationFlag = "Hangar"
				) S ON S.characterID = T.tokenCharacterID
			WHERE lastStatus != "invalid_token" AND scopes LIKE "%read_structures%"
			`)
	if err != nil {
		return err
	}
	defer structures.Close()

	work := []redisqueue.Work{}

	// Loop the structure pairs
	for structures.Next() {
		var (
			characterID, tokenCharacterID int32
			structureID                   int64
		)

		err = structures.Scan(&characterID, &tokenCharacterID, &structureID)
		if err != nil {
			return err
		}
		if !s.inQueue.CheckWorkExpired("evedata_structurechar_failure",
			fmt.Sprintf("%d%d", structureID, tokenCharacterID)) {
			work = append(work, redisqueue.Work{Operation: "characterStructures",
				Parameter: []interface{}{
					characterID,
					tokenCharacterID,
					structureID,
				}})
		}
	}

	return s.QueueWork(work, redisqueue.Priority_Lowest)
}

func characterStructureMarket(s *Artifice) error {
	structures, err := s.db.Query(
		`SELECT T.characterID, T.tokenCharacterID, S.locationID
			FROM evedata.crestTokens T
			INNER JOIN (
					SELECT DISTINCT characterID, locationID FROM evedata.orders WHERE locationID > 70000000
					UNION DISTINCT
					SELECT DISTINCT characterID, locationID FROM evedata.assets WHERE locationID > 70000000
						AND locationFlag = "Hangar"
				) S ON S.characterID = T.tokenCharacterID
			WHERE lastStatus != "invalid_token" AND scopes LIKE "%read_structures%"
			`)
	if err != nil {
		return err
	}
	defer structures.Close()

	work := []redisqueue.Work{}

	// Loop the structure pairs
	for structures.Next() {
		var (
			characterID, tokenCharacterID int32
			structureID                   int64
		)

		err = structures.Scan(&characterID, &tokenCharacterID, &structureID)
		if err != nil {
			return err
		}

		if !s.inQueue.CheckWorkExpired("evedata_structuremarket_failure",
			fmt.Sprintf("%d%d", structureID, tokenCharacterID)) {
			work = append(work, redisqueue.Work{Operation: "characterStructureMarket",
				Parameter: []interface{}{
					characterID,
					tokenCharacterID,
					structureID,
				}})
			log.Printf("adding %d\n", structureID)
		} else {
			log.Printf("failed, ignoring %d\n", structureID)
		}
	}

	return s.QueueWork(work, redisqueue.Priority_Lowest)
}

func characterOrders(s *Artifice) error {
	work := []redisqueue.Work{}
	if pairs, err := s.GetCharactersForScope("read_character_orders"); err != nil {
		return err
	} else {
		for _, p := range pairs {
			work = append(work, redisqueue.Work{Operation: "characterOrders", Parameter: []int32{p.CharacterID, p.TokenCharacterID}})
		}
	}

	return s.QueueWork(work, redisqueue.Priority_High)
}

func characterNotifications(s *Artifice) error {
	work := []redisqueue.Work{}
	if pairs, err := s.GetCharactersForScope("read_notifications"); err != nil {
		return err
	} else {
		for _, p := range pairs {
			work = append(work, redisqueue.Work{Operation: "characterNotifications", Parameter: []int32{p.CharacterID, p.TokenCharacterID}})
		}
	}

	return s.QueueWork(work, redisqueue.Priority_High)
}

func allianceContacts(s *Artifice) error {
	work := []redisqueue.Work{}
	if pairs, err := s.GetAllianceForScope("alliances.read_contacts"); err != nil {
		return err
	} else {
		for _, p := range pairs {
			work = append(work, redisqueue.Work{Operation: "allianceContacts", Parameter: []int32{p.CharacterID, p.TokenCharacterID, p.AllianceID}})
		}
	}

	return s.QueueWork(work, redisqueue.Priority_High)
}

func corporationContacts(s *Artifice) error {
	work := []redisqueue.Work{}
	if pairs, err := s.GetCorporationForScope("corporations.read_contacts"); err != nil {
		return err
	} else {
		for _, p := range pairs {
			work = append(work, redisqueue.Work{Operation: "corporationContacts", Parameter: []int32{p.CharacterID, p.TokenCharacterID, p.CorporationID}})
		}
	}

	return s.QueueWork(work, redisqueue.Priority_High)
}

func characterContactSync(s *Artifice) error {
	entities, err := s.db.Query(
		`SELECT S.characterID, source, group_concat(destination) AS destinations
		FROM evedata.contactSyncs S
		INNER JOIN evedata.crestTokens T ON T.tokenCharacterID = destination
		WHERE lastStatus != "invalid_token" AND scopes LIKE ?
		GROUP BY source
		HAVING max(nextSync) < UTC_TIMESTAMP();`, "%characters.read_contacts%")
	if err != nil {
		return err
	}
	defer entities.Close()

	work := []redisqueue.Work{}

	// Loop the entities
	for entities.Next() {
		var (
			cid, source  int32
			destinations string
		)

		err = entities.Scan(&cid, &source, &destinations)
		if err != nil {
			return err
		}

		work = append(work, redisqueue.Work{Operation: "characterContactSync", Parameter: []interface{}{
			cid,
			source,
			destinations,
		}})
	}

	return s.QueueWork(work, redisqueue.Priority_Normal)
}

func characterAuthOwners(s *Artifice) error {
	work := []redisqueue.Work{}
	if pairs, err := s.GetCharactersForScope("read_corporation_roles"); err != nil {
		return err
	} else {
		for _, p := range pairs {
			work = append(work, redisqueue.Work{Operation: "characterAuthOwner", Parameter: []int32{p.CharacterID, p.TokenCharacterID}})
		}
	}

	return s.QueueWork(work, redisqueue.Priority_High)
}

// This is sensitive so we will do it here to prevent mixing it with public data.
// figure out character alliance and corp for our members
func crestCharacters(s *Artifice) error {
	var chars []int32
	err := s.db.Select(&chars,
		`SELECT DISTINCT tokenCharacterID FROM evedata.crestTokens`)
	if err != nil {
		log.Println(err)
		return err
	}

	// Get a list of characters sharing data to check for changes
	sharing := make(map[int32]int32)
	rows, err := s.db.Query(`
		SELECT DISTINCT T.tokenCharacterID, T.corporationID FROM evedata.crestTokens T
			INNER JOIN evedata.sharing S ON S.tokenCharacterID = T.tokenCharacterID`)
	if err != nil {
		log.Println(err)
		return err
	}
	defer rows.Close()
	for rows.Next() {
		var char, corp int32
		err := rows.Scan(&char, &corp)
		if err != nil {
			log.Println(err)
			return err
		}
		sharing[char] = corp
	}

	for start := 0; start < len(chars); start = start + 50 {
		end := min(start+50, len(chars))
		if affiliation, _, err := s.esi.ESI.CharacterApi.PostCharactersAffiliation(context.Background(), chars[start:end], nil); err != nil {
			log.Println(err)
			continue
		} else {
			tx, err := s.db.Beginx()
			defer tx.Rollback()
			if err != nil {
				return err
			}
			for _, c := range affiliation {
				// See if they changed corporation, if they have shares, warn them they are still sharing.
				if check, ok := sharing[c.CharacterId]; ok {
					if check != c.CorporationId {
						s.mailCorporationChangeWithShares(c.CharacterId)
					}
				}

				if _, err := tx.Exec(`UPDATE evedata.crestTokens
					SET corporationID = ?, allianceID = ?, factionID = ?
					WHERE tokenCharacterID = ?;
					`, c.CorporationId, c.AllianceId, c.FactionId, c.CharacterId); err != nil {
					log.Println(err)
					continue
				}
			}
			if err = sqlhelper.RetryTransaction(tx); err != nil {
				return err
			}
		}

		// Update integrations services with factionID.
		if err = s.doSQL(`
			UPDATE evedata.integrations S
			INNER JOIN evedata.crestTokens C ON C.allianceID = S.entityID OR C.corporationID = S.entityID
				AND expiry > DATE_SUB(UTC_TIMESTAMP(), INTERVAL 20 MINUTE)
			SET S.factionID = C.factionID`); err != nil {
			return err
		}
	}
	return nil
}

func min(x, y int) int {
	if x < y {
		return x
	}
	return y
}
<commit_msg>dont be so verbose<commit_after>package artifice

import (
	"context"
	"fmt"
	"log"
	"time"

	"github.com/antihax/evedata/internal/redisqueue"
	"github.com/antihax/evedata/internal/sqlhelper"
)

func init() {
	registerTrigger("characterTransactions", characterTransactions, time.NewTicker(time.Second*3600))
	registerTrigger("characterAssets", characterAssets, time.NewTicker(time.Second*3600))
	registerTrigger("characterOrders", characterOrders, time.NewTicker(time.Second*1200))
	registerTrigger("characterStructures", characterStructures, time.NewTicker(time.Second*3600))
	registerTrigger("characterStructureMarket", characterStructureMarket, time.NewTicker(time.Second*300))

	registerTrigger("characterNotifications", characterNotifications, time.NewTicker(time.Second*600))
	registerTrigger("characterContactSync", characterContactSync, time.NewTicker(time.Second*360))
	registerTrigger("characterAuthOwners", characterAuthOwners, time.NewTicker(time.Second*3600))
	registerTrigger("crestCharacters", crestCharacters, time.NewTicker(time.Second*3600))

	registerTrigger("allianceContacts", allianceContacts, time.NewTicker(time.Second*360))
	registerTrigger("corporationContacts", corporationContacts, time.NewTicker(time.Second*360))
}

func characterTransactions(s *Artifice) error {
	if pairs, err := s.GetCharactersForScope("read_character_wallet"); err != nil {
		return err
	} else {
		work := []redisqueue.Work{}
		for _, p := range pairs {
			work = append(work, redisqueue.Work{Operation: "characterWalletTransactions", Parameter: []int32{p.CharacterID, p.TokenCharacterID}})
			work = append(work, redisqueue.Work{Operation: "characterWalletJournal", Parameter: []int32{p.CharacterID, p.TokenCharacterID}})
		}
		return s.QueueWork(work, redisqueue.Priority_Normal)
	}
}

func characterAssets(s *Artifice) error {
	work := []redisqueue.Work{}
	if pairs, err := s.GetCharactersForScope("read_assets"); err != nil {
		return err
	} else {
		for _, p := range pairs {
			work = append(work, redisqueue.Work{Operation: "characterAssets", Parameter: []int32{p.CharacterID, p.TokenCharacterID}})
		}
	}

	return s.QueueWork(work, redisqueue.Priority_High)
}

func characterStructures(s *Artifice) error {
	structures, err := s.db.Query(
		`SELECT T.characterID, T.tokenCharacterID, S.locationID
			FROM evedata.crestTokens T
			INNER JOIN (
					SELECT DISTINCT characterID, locationID FROM evedata.orders WHERE locationID > 70000000
					UNION DISTINCT
					SELECT DISTINCT characterID, locationID FROM evedata.assets WHERE locationID > 70000000
						AND locationFlag = "Hangar"
				) S ON S.characterID = T.tokenCharacterID
			WHERE lastStatus != "invalid_token" AND scopes LIKE "%read_structures%"
			`)
	if err != nil {
		return err
	}
	defer structures.Close()

	work := []redisqueue.Work{}

	// Loop the structure pairs
	for structures.Next() {
		var (
			characterID, tokenCharacterID int32
			structureID                   int64
		)

		err = structures.Scan(&characterID, &tokenCharacterID, &structureID)
		if err != nil {
			return err
		}
		if !s.inQueue.CheckWorkExpired("evedata_structurechar_failure",
			fmt.Sprintf("%d%d", structureID, tokenCharacterID)) {
			work = append(work, redisqueue.Work{Operation: "characterStructures",
				Parameter: []interface{}{
					characterID,
					tokenCharacterID,
					structureID,
				}})
		}
	}

	return s.QueueWork(work, redisqueue.Priority_Lowest)
}

func characterStructureMarket(s *Artifice) error {
	structures, err := s.db.Query(
		`SELECT T.characterID, T.tokenCharacterID, S.locationID
			FROM evedata.crestTokens T
			INNER JOIN (
					SELECT DISTINCT characterID, locationID FROM evedata.orders WHERE locationID > 70000000
					UNION DISTINCT
					SELECT DISTINCT characterID, locationID FROM evedata.assets WHERE locationID > 70000000
						AND locationFlag = "Hangar"
				) S ON S.characterID = T.tokenCharacterID
			WHERE lastStatus != "invalid_token" AND scopes LIKE "%read_structures%"
			`)
	if err != nil {
		return err
	}
	defer structures.Close()

	work := []redisqueue.Work{}

	// Loop the structure pairs
	for structures.Next() {
		var (
			characterID, tokenCharacterID int32
			structureID                   int64
		)

		err = structures.Scan(&characterID, &tokenCharacterID, &structureID)
		if err != nil {
			return err
		}

		if !s.inQueue.CheckWorkExpired("evedata_structuremarket_failure",
			fmt.Sprintf("%d%d", structureID, tokenCharacterID)) {
			work = append(work, redisqueue.Work{Operation: "characterStructureMarket",
				Parameter: []interface{}{
					characterID,
					tokenCharacterID,
					structureID,
				}})
		}
	}

	return s.QueueWork(work, redisqueue.Priority_Lowest)
}

func characterOrders(s *Artifice) error {
	work := []redisqueue.Work{}
	if pairs, err := s.GetCharactersForScope("read_character_orders"); err != nil {
		return err
	} else {
		for _, p := range pairs {
			work = append(work, redisqueue.Work{Operation: "characterOrders", Parameter: []int32{p.CharacterID, p.TokenCharacterID}})
		}
	}

	return s.QueueWork(work, redisqueue.Priority_High)
}

func characterNotifications(s *Artifice) error {
	work := []redisqueue.Work{}
	if pairs, err := s.GetCharactersForScope("read_notifications"); err != nil {
		return err
	} else {
		for _, p := range pairs {
			work = append(work, redisqueue.Work{Operation: "characterNotifications", Parameter: []int32{p.CharacterID, p.TokenCharacterID}})
		}
	}

	return s.QueueWork(work, redisqueue.Priority_High)
}

func allianceContacts(s *Artifice) error {
	work := []redisqueue.Work{}
	if pairs, err := s.GetAllianceForScope("alliances.read_contacts"); err != nil {
		return err
	} else {
		for _, p := range pairs {
			work = append(work, redisqueue.Work{Operation: "allianceContacts", Parameter: []int32{p.CharacterID, p.TokenCharacterID, p.AllianceID}})
		}
	}

	return s.QueueWork(work, redisqueue.Priority_High)
}

func corporationContacts(s *Artifice) error {
	work := []redisqueue.Work{}
	if pairs, err := s.GetCorporationForScope("corporations.read_contacts"); err != nil {
		return err
	} else {
		for _, p := range pairs {
			work = append(work, redisqueue.Work{Operation: "corporationContacts", Parameter: []int32{p.CharacterID, p.TokenCharacterID, p.CorporationID}})
		}
	}

	return s.QueueWork(work, redisqueue.Priority_High)
}

func characterContactSync(s *Artifice) error {
	entities, err := s.db.Query(
		`SELECT S.characterID, source, group_concat(destination) AS destinations
		FROM evedata.contactSyncs S
		INNER JOIN evedata.crestTokens T ON T.tokenCharacterID = destination
		WHERE lastStatus != "invalid_token" AND scopes LIKE ?
		GROUP BY source
		HAVING max(nextSync) < UTC_TIMESTAMP();`, "%characters.read_contacts%")
	if err != nil {
		return err
	}
	defer entities.Close()

	work := []redisqueue.Work{}

	// Loop the entities
	for entities.Next() {
		var (
			cid, source  int32
			destinations string
		)

		err = entities.Scan(&cid, &source, &destinations)
		if err != nil {
			return err
		}

		work = append(work, redisqueue.Work{Operation: "characterContactSync", Parameter: []interface{}{
			cid,
			source,
			destinations,
		}})
	}

	return s.QueueWork(work, redisqueue.Priority_Normal)
}

func characterAuthOwners(s *Artifice) error {
	work := []redisqueue.Work{}
	if pairs, err := s.GetCharactersForScope("read_corporation_roles"); err != nil {
		return err
	} else {
		for _, p := range pairs {
			work = append(work, redisqueue.Work{Operation: "characterAuthOwner", Parameter: []int32{p.CharacterID, p.TokenCharacterID}})
		}
	}

	return s.QueueWork(work, redisqueue.Priority_High)
}

// This is sensitive so we will do it here to prevent mixing it with public data.
// figure out character alliance and corp for our members
func crestCharacters(s *Artifice) error {
	var chars []int32
	err := s.db.Select(&chars,
		`SELECT DISTINCT tokenCharacterID FROM evedata.crestTokens`)
	if err != nil {
		log.Println(err)
		return err
	}

	// Get a list of characters sharing data to check for changes
	sharing := make(map[int32]int32)
	rows, err := s.db.Query(`
		SELECT DISTINCT T.tokenCharacterID, T.corporationID FROM evedata.crestTokens T
			INNER JOIN evedata.sharing S ON S.tokenCharacterID = T.tokenCharacterID`)
	if err != nil {
		log.Println(err)
		return err
	}
	defer rows.Close()
	for rows.Next() {
		var char, corp int32
		err := rows.Scan(&char, &corp)
		if err != nil {
			log.Println(err)
			return err
		}
		sharing[char] = corp
	}

	for start := 0; start < len(chars); start = start + 50 {
		end := min(start+50, len(chars))
		if affiliation, _, err := s.esi.ESI.CharacterApi.PostCharactersAffiliation(context.Background(), chars[start:end], nil); err != nil {
			log.Println(err)
			continue
		} else {
			tx, err := s.db.Beginx()
			defer tx.Rollback()
			if err != nil {
				return err
			}
			for _, c := range affiliation {
				// See if they changed corporation, if they have shares, warn them they are still sharing.
				if check, ok := sharing[c.CharacterId]; ok {
					if check != c.CorporationId {
						s.mailCorporationChangeWithShares(c.CharacterId)
					}
				}

				if _, err := tx.Exec(`UPDATE evedata.crestTokens
					SET corporationID = ?, allianceID = ?, factionID = ?
					WHERE tokenCharacterID = ?;
					`, c.CorporationId, c.AllianceId, c.FactionId, c.CharacterId); err != nil {
					log.Println(err)
					continue
				}
			}
			if err = sqlhelper.RetryTransaction(tx); err != nil {
				return err
			}
		}

		// Update integrations services with factionID.
		if err = s.doSQL(`
			UPDATE evedata.integrations S
			INNER JOIN evedata.crestTokens C ON C.allianceID = S.entityID OR C.corporationID = S.entityID
				AND expiry > DATE_SUB(UTC_TIMESTAMP(), INTERVAL 20 MINUTE)
			SET S.factionID = C.factionID`); err != nil {
			return err
		}
	}
	return nil
}

func min(x, y int) int {
	if x < y {
		return x
	}
	return y
}
<|endoftext|>


<commit_before>package websession

import (
	"OnlineJudge/base"
	locals "OnlineJudge/sessions"

	"github.com/gorilla/sessions"

	"errors"
	"net/http"
)

var (
	KeyUsername  = ".username"
	KeyPrivilege = ".privilege"
	KeyUserId    = ".user_id"
	KeyFlash     = ".flash"
	KeyIPAddr    = ".ipaddr"

	ErrNotFound   = errors.New("Key not in the session")
	ErrKeyInvalid = errors.New("Key invalid")
)

type WebSession struct {
	session *sessions.Session
}

func (this *WebSession) IsLogin() bool {
	if name, ok := this.session.Values[KeyUsername].(string); ok {
		if name != "" {
			return true
		}
	}
	return false
}

func (this *WebSession) IsRoot() bool {
	if privilege, ok := this.session.Values[KeyPrivilege].(string); ok {
		if privilege == "root" {
			return true
		}
	}
	return false
}

func (this *WebSession) GetUsername() string {
	username, ok := this.session.Values[KeyUsername].(string)
	if !ok {
		return ""
	}
	return username
}

func (this *WebSession) GetPrivilege() string {
	privilege, ok := this.session.Values[KeyPrivilege].(string)
	if !ok {
		return ""
	}
	return privilege
}

func (this *WebSession) GetUserId() int64 {
	userid, ok := this.session.Values[KeyUserId].(int64)
	if !ok {
		return 0
	}
	return userid
}

func (this *WebSession) GetIPAddr() string {
	ip, ok := this.session.Values[KeyIPAddr].(string)
	if !ok {
		return ""
	}
	return ip
}

func (this *WebSession) Get(key interface{}) (interface{}, error) {
	if _, ok := this.session.Values[key]; !ok {
		return nil, ErrNotFound
	}
	return this.session.Values[key], nil
}

func (this *WebSession) Set(key interface{}, val interface{}) error {
	if val, ok := key.(string); ok {
		if !base.IsNilOrZero(val) {
			if val[0] == '.' {
				return ErrKeyInvalid
			}
		}
	}
	this.session.Values[key] = val
	return nil
}

func (this *WebSession) SetUsername(username string) {
	this.session.Values[KeyUsername] = username
}

func (this *WebSession) SetPrivilege(privilege string) {
	this.session.Values[KeyPrivilege] = privilege
}

func (this *WebSession) SetUserId(user_id int64) {
	this.session.Values[KeyUserId] = user_id
}

func (this *WebSession) SetIPAddr(ip string) {
	this.session.Values[KeyIPAddr] = ip
}

func (this *WebSession) Flashes(vars ...string) []interface{} {
	return this.session.Flashes(vars...)
}

func (this *WebSession) AddFlash(value interface{}, vars ...string) {
	this.session.AddFlash(value, vars...)
}

func (this *WebSession) Save(h *http.Request, w http.ResponseWriter) {
	this.session.Save(h, w)
}

func (this *WebSession) Logout() {
	for k, _ := range this.session.Values {
		delete(this.session.Values, k)
	}
}

func NewSession(sess *sessions.Session) locals.Session {
	newsess := &WebSession{
		session: sess,
	}
	var session locals.Session = newsess
	return session
}

func NewWebSession(sess *sessions.Session) *WebSession {
	return &WebSession{
		session: sess,
	}
}
<commit_msg>right way to delete a session<commit_after>package websession

import (
	"OnlineJudge/base"
	locals "OnlineJudge/sessions"

	"github.com/gorilla/sessions"

	"errors"
	"net/http"
)

var (
	KeyUsername  = ".username"
	KeyPrivilege = ".privilege"
	KeyUserId    = ".user_id"
	KeyFlash     = ".flash"
	KeyIPAddr    = ".ipaddr"

	ErrNotFound   = errors.New("Key not in the session")
	ErrKeyInvalid = errors.New("Key invalid")
)

type WebSession struct {
	session *sessions.Session
}

func (this *WebSession) IsLogin() bool {
	if name, ok := this.session.Values[KeyUsername].(string); ok {
		if name != "" {
			return true
		}
	}
	return false
}

func (this *WebSession) IsRoot() bool {
	if privilege, ok := this.session.Values[KeyPrivilege].(string); ok {
		if privilege == "root" {
			return true
		}
	}
	return false
}

func (this *WebSession) GetUsername() string {
	username, ok := this.session.Values[KeyUsername].(string)
	if !ok {
		return ""
	}
	return username
}

func (this *WebSession) GetPrivilege() string {
	privilege, ok := this.session.Values[KeyPrivilege].(string)
	if !ok {
		return ""
	}
	return privilege
}

func (this *WebSession) GetUserId() int64 {
	userid, ok := this.session.Values[KeyUserId].(int64)
	if !ok {
		return 0
	}
	return userid
}

func (this *WebSession) GetIPAddr() string {
	ip, ok := this.session.Values[KeyIPAddr].(string)
	if !ok {
		return ""
	}
	return ip
}

func (this *WebSession) Get(key interface{}) (interface{}, error) {
	if _, ok := this.session.Values[key]; !ok {
		return nil, ErrNotFound
	}
	return this.session.Values[key], nil
}

func (this *WebSession) Set(key interface{}, val interface{}) error {
	if val, ok := key.(string); ok {
		if !base.IsNilOrZero(val) {
			if val[0] == '.' {
				return ErrKeyInvalid
			}
		}
	}
	this.session.Values[key] = val
	return nil
}

func (this *WebSession) SetUsername(username string) {
	this.session.Values[KeyUsername] = username
}

func (this *WebSession) SetPrivilege(privilege string) {
	this.session.Values[KeyPrivilege] = privilege
}

func (this *WebSession) SetUserId(user_id int64) {
	this.session.Values[KeyUserId] = user_id
}

func (this *WebSession) SetIPAddr(ip string) {
	this.session.Values[KeyIPAddr] = ip
}

func (this *WebSession) Flashes(vars ...string) []interface{} {
	return this.session.Flashes(vars...)
}

func (this *WebSession) AddFlash(value interface{}, vars ...string) {
	this.session.AddFlash(value, vars...)
}

func (this *WebSession) Save(h *http.Request, w http.ResponseWriter) {
	this.session.Save(h, w)
}

func (this *WebSession) Logout() {
	this.session.Options.MaxAge = -1
}

func NewSession(sess *sessions.Session) locals.Session {
	newsess := &WebSession{
		session: sess,
	}
	var session locals.Session = newsess
	return session
}

func NewWebSession(sess *sessions.Session) *WebSession {
	return &WebSession{
		session: sess,
	}
}
<|endoftext|>


<commit_before>package api

// InstanceConsoleControl represents a message on the instance console "control" socket.
//
// API extension: instances
type InstanceConsoleControl struct {
	Command string            `json:"command" yaml:"command"`
	Args    map[string]string `json:"args" yaml:"args"`
}

// InstanceConsolePost represents a LXD instance console request.
//
// API extension: instances
type InstanceConsolePost struct {
	Width  int `json:"width" yaml:"width"`
	Height int `json:"height" yaml:"height"`
}
<commit_msg>shared/api: Add Type field to InstanceConsolePost<commit_after>package api

// InstanceConsoleControl represents a message on the instance console "control" socket.
//
// API extension: instances
type InstanceConsoleControl struct {
	Command string            `json:"command" yaml:"command"`
	Args    map[string]string `json:"args" yaml:"args"`
}

// InstanceConsolePost represents a LXD instance console request.
//
// API extension: instances
type InstanceConsolePost struct {
	Width  int `json:"width" yaml:"width"`
	Height int `json:"height" yaml:"height"`

	// API extension: console_vga_type
	Type string `json:"type" yaml:"type"`
}
<|endoftext|>


<commit_before>package engine

import (
	"io"

	"github.com/keybase/client/go/libkb"
	"golang.org/x/crypto/openpgp"
)

type PGPEncryptArg struct {
	Source     io.Reader
	Sink       io.Writer
	Signer     *libkb.PgpKeyBundle
	Recipients []*libkb.PgpKeyBundle
	Sign       bool
}

type PGPEncrypt struct {
	arg *PGPEncryptArg
}

// NewPGPEncrypt creates a PGPEncrypt engine.
func NewPGPEncrypt(arg *PGPEncryptArg) *PGPEncrypt {
	return &PGPEncrypt{arg: arg}
}

// Name is the unique engine name.
func (e *PGPEncrypt) Name() string {
	return "PGPEncrypt"
}

// GetPrereqs returns the engine prereqs.
func (e *PGPEncrypt) GetPrereqs() EnginePrereqs {
	return EnginePrereqs{}
}

// RequiredUIs returns the required UIs.
func (e *PGPEncrypt) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{}
}

// SubConsumers returns the other UI consumers for this engine.
func (e *PGPEncrypt) SubConsumers() []libkb.UIConsumer {
	return nil
}

// Run starts the engine.
func (e *PGPEncrypt) Run(ctx *Context, args, reply interface{}) error {
	to := make([]*openpgp.Entity, len(e.arg.Recipients))
	for i, r := range e.arg.Recipients {
		to[i] = (*openpgp.Entity)(r)
	}
	w, err := openpgp.Encrypt(e.arg.Sink, to, (*openpgp.Entity)(e.arg.Signer), nil, nil)
	if err != nil {
		return err
	}
	n, err := io.Copy(w, e.arg.Source)
	if err != nil {
		return err
	}
	G.Log.Debug("PGPEncrypt.Run: wrote %d bytes", n)
	return w.Close()
}
<commit_msg>documentation<commit_after>package engine

import (
	"io"

	"github.com/keybase/client/go/libkb"
	"golang.org/x/crypto/openpgp"
)

// PGPEncryptArg contains the arguments for the PGPEncrypt engine.
// Source is the source reader, Sink is the destination writer.
// Signer is optional.  If provided, it will be used to sign the
// message.  Recipients is a list of PGP keys for the recipients
// of this message.
type PGPEncryptArg struct {
	Source     io.Reader
	Sink       io.Writer
	Signer     *libkb.PgpKeyBundle
	Recipients []*libkb.PgpKeyBundle
}

// PGPEncrypt is an engine for pgp encrypting data to one or
// multiple recipients, optionally signed.
type PGPEncrypt struct {
	arg *PGPEncryptArg
}

// NewPGPEncrypt creates a PGPEncrypt engine.
func NewPGPEncrypt(arg *PGPEncryptArg) *PGPEncrypt {
	return &PGPEncrypt{arg: arg}
}

// Name is the unique engine name.
func (e *PGPEncrypt) Name() string {
	return "PGPEncrypt"
}

// GetPrereqs returns the engine prereqs.
func (e *PGPEncrypt) GetPrereqs() EnginePrereqs {
	return EnginePrereqs{}
}

// RequiredUIs returns the required UIs.
func (e *PGPEncrypt) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{}
}

// SubConsumers returns the other UI consumers for this engine.
func (e *PGPEncrypt) SubConsumers() []libkb.UIConsumer {
	return nil
}

// Run starts the engine.
func (e *PGPEncrypt) Run(ctx *Context, args, reply interface{}) error {
	to := make([]*openpgp.Entity, len(e.arg.Recipients))
	for i, r := range e.arg.Recipients {
		to[i] = (*openpgp.Entity)(r)
	}
	w, err := openpgp.Encrypt(e.arg.Sink, to, (*openpgp.Entity)(e.arg.Signer), nil, nil)
	if err != nil {
		return err
	}
	n, err := io.Copy(w, e.arg.Source)
	if err != nil {
		return err
	}
	G.Log.Debug("PGPEncrypt.Run: wrote %d bytes", n)
	return w.Close()
}
<|endoftext|>


<commit_before>package fixchain

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"time"
)

type urlCache struct {
	client *http.Client
	cache  map[string][]byte
	// counters may not be totally accurate due to non-atomicity
	hit       uint
	miss      uint
	errors    uint
	badStatus uint
	readFail  uint
}

func (u *urlCache) getURL(url string) ([]byte, error) {
	r, ok := u.cache[url]
	if ok {
		u.hit++
		return r, nil
	}
	c, err := u.client.Get(url)
	if err != nil {
		u.errors++
		return nil, err
	}
	defer c.Body.Close()
	if c.StatusCode != 200 {
		u.badStatus++
		return nil, fmt.Errorf("can't deal with status %d", c.StatusCode)
	}
	r, err = ioutil.ReadAll(c.Body)
	if err != nil {
		u.readFail++
		return nil, err
	}
	u.miss++
	u.cache[url] = r
	return r, nil
}

func newURLCache(c *http.Client, logStats bool) *urlCache {
	u := &urlCache{cache: make(map[string][]byte), client: c}

	if logStats {
		t := time.NewTicker(time.Second)
		go func() {
			for _ = range t.C {
				log.Printf("cache: %d hits, %d misses, %d errors, "+
					"%d bad status, %d read fail, %d cached", u.hit,
					u.miss, u.errors, u.badStatus, u.readFail,
					len(u.cache))
			}
		}()
	}

	return u
}
<commit_msg>Add reminder to do permanent http error caching<commit_after>package fixchain

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"time"
)

type urlCache struct {
	client *http.Client
	cache  map[string][]byte
	// counters may not be totally accurate due to non-atomicity
	hit       uint
	miss      uint
	errors    uint
	badStatus uint
	readFail  uint
}

func (u *urlCache) getURL(url string) ([]byte, error) {
	r, ok := u.cache[url]
	if ok {
		u.hit++
		return r, nil
	}
	c, err := u.client.Get(url)
	if err != nil {
		u.errors++
		return nil, err
	}
	defer c.Body.Close()
	// TODO(katjoyce): Add caching of permanent errors.
	if c.StatusCode != 200 {
		u.badStatus++
		return nil, fmt.Errorf("can't deal with status %d", c.StatusCode)
	}
	r, err = ioutil.ReadAll(c.Body)
	if err != nil {
		u.readFail++
		return nil, err
	}
	u.miss++
	u.cache[url] = r
	return r, nil
}

func newURLCache(c *http.Client, logStats bool) *urlCache {
	u := &urlCache{cache: make(map[string][]byte), client: c}

	if logStats {
		t := time.NewTicker(time.Second)
		go func() {
			for _ = range t.C {
				log.Printf("cache: %d hits, %d misses, %d errors, "+
					"%d bad status, %d read fail, %d cached", u.hit,
					u.miss, u.errors, u.badStatus, u.readFail,
					len(u.cache))
			}
		}()
	}

	return u
}
<|endoftext|>


<commit_before>package main

// Script Functionality

import (
	"fmt"
	"github.com/spf13/cobra"
	"github.com/stroblindustries/coreutils"
	"os"
	"path/filepath"
	"strings"
)

var scriptCmd = &cobra.Command{
	Use:     "script",
	Aliases: []string{"run-script"},
	Short:   "Run a custom script",
	Long:    "Run a custom script",
	Run:     script,
}

var verbose bool

func init() {
	scriptCmd.Flags().BoolVarP(&verbose, "verbose", "v", false, "Enable verbose mode.")
}

func script(cmd *cobra.Command, args []string) {
	if len(args) != 0 { // If no argument are passed
		for _, arg := range args { // For each argument
			if _, exists := noodles.Scripts[arg]; exists { // If this
				RunScript(arg)
			} else {
				fmt.Printf("%s is not a valid script. Exiting.\n", arg)
				break
			}
		}
	} else {
		if len(noodles.Scripts) != 0 { // If there are scripts set
			fmt.Printf("No scripts passed to run. Here is a list:\n\n")
			for name, script := range noodles.Scripts {
				line := name

				if script.Description != "" {
					line += ": " + script.Description
				}

				fmt.Println("-", line)
			}
			fmt.Printf("\n") // Intentional padding ad end of output
		} else {
			fmt.Println("No scripts defined in your noodles.toml file.")
		}
	}
}

// RunScript will run the script provided
func RunScript(name string) {
	script, _ := noodles.Scripts[name] // Get our script

	if script.Exec != "" { // If there is an executable
		if script.Directory != "" { // If we should run this command in a directory
			failedToChange := os.Chdir(filepath.Join(workdir, script.Directory)) // Change to the directory

			if failedToChange != nil { // If we failed to change to the directory
				fmt.Printf("Failed to change to the following directory: %s\n", script.Directory)

				if verbose {
					fmt.Printf("Full error: %s\n", failedToChange)
				}

				return // Don't continue with exec
			}
		}

		if verbose {
			commandRunning := script.Exec + " " + (strings.Join(script.Arguments, " "))
			fmt.Printf("Running: %s\n", commandRunning)
		}

		output := coreutils.ExecCommand(script.Exec, script.Arguments, script.Redirect)

		if (script.File != "") && script.Redirect { // If we should redirect output to a file
			coreutils.WriteOrUpdateFile(script.File, []byte(output), coreutils.NonGlobalFileMode)
		}

		os.Chdir(workdir) // Change back to the work dir if needed
	} else {
		fmt.Printf("No executable set for the script: %s", name)
	}
}
<commit_msg>Simplify script primary func and support -s / --script arg.<commit_after>package main

// Script Functionality

import (
	"fmt"
	"github.com/spf13/cobra"
	"github.com/stroblindustries/coreutils"
	"os"
	"path/filepath"
	"strings"
)

var scriptCmd = &cobra.Command{
	Use:     "script",
	Aliases: []string{"run-script"},
	Short:   "Run a custom script",
	Long:    "Run a custom script",
	Run:     script,
}

var verbose bool
var selectedScript string

func init() {
	scriptCmd.Flags().BoolVarP(&verbose, "verbose", "v", false, "Enable verbose mode.")
	scriptCmd.Flags().StringVarP(&selectedScript, "script", "s", "", "Name of the script we're running")
}

func script(cmd *cobra.Command, args []string) {
	if selectedScript == "" { // If no script is set
		for name := range noodles.Scripts {
			RunScript(name)
		}
	} else { // If a script is set
		RunScript(selectedScript)
	}
}

// RunScript will run the script provided
func RunScript(name string) {
	script, _ := noodles.Scripts[name] // Get our script

	if script.Exec != "" { // If there is an executable
		if script.Directory != "" { // If we should run this command in a directory
			failedToChange := os.Chdir(filepath.Join(workdir, script.Directory)) // Change to the directory

			if failedToChange != nil { // If we failed to change to the directory
				fmt.Printf("Failed to change to the following directory: %s\n", script.Directory)

				if verbose {
					fmt.Printf("Full error: %s\n", failedToChange)
				}

				return // Don't continue with exec
			}
		}

		if verbose {
			commandRunning := script.Exec + " " + (strings.Join(script.Arguments, " "))
			fmt.Printf("Running: %s\n", commandRunning)
		}

		output := coreutils.ExecCommand(script.Exec, script.Arguments, script.Redirect)

		if (script.File != "") && script.Redirect { // If we should redirect output to a file
			coreutils.WriteOrUpdateFile(script.File, []byte(output), coreutils.NonGlobalFileMode)
		}

		os.Chdir(workdir) // Change back to the work dir if needed
	} else {
		fmt.Printf("No executable set for the script: %s", name)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package worker

import (
	"errors"
	"fmt"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/vt/logutil"
	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/wrangler"
	"golang.org/x/net/context"
)

// Instance encapsulate the execution state of vtworker.
type Instance struct {
	// Default wrangler for all operations.
	// Users can specify their own in RunCommand() e.g. the gRPC server does this.
	wr *wrangler.Wrangler

	// mutex is protecting all the following variables
	// 3 states here:
	// - no job ever ran (or reset was run): currentWorker is nil,
	// currentContext/currentCancelFunc is nil, lastRunError is nil
	// - one worker running: currentWorker is set,
	//   currentContext/currentCancelFunc is set, lastRunError is nil
	// - (at least) one worker already ran, none is running atm:
	//   currentWorker is set, currentContext is nil, lastRunError
	//   has the error returned by the worker.
	currentWorkerMutex  sync.Mutex
	currentWorker       Worker
	currentMemoryLogger *logutil.MemoryLogger
	currentContext      context.Context
	currentCancelFunc   context.CancelFunc
	lastRunError        error

	// backgroundContext is context.Background() from main() which has to be plumbed through.
	backgroundContext      context.Context
	topoServer             topo.Server
	cell                   string
	commandDisplayInterval time.Duration
}

// NewInstance creates a new Instance.
func NewInstance(ctx context.Context, ts topo.Server, cell string, commandDisplayInterval time.Duration) *Instance {
	wi := &Instance{backgroundContext: ctx, topoServer: ts, cell: cell, commandDisplayInterval: commandDisplayInterval}
	// Note: setAndStartWorker() also adds a MemoryLogger for the webserver.
	wi.wr = wi.CreateWrangler(logutil.NewConsoleLogger())
	return wi
}

// CreateWrangler creates a new wrangler using the instance specific configuration.
func (wi *Instance) CreateWrangler(logger logutil.Logger) *wrangler.Wrangler {
	return wrangler.New(logger, wi.topoServer, tmclient.NewTabletManagerClient())
}

// setAndStartWorker will set the current worker.
// We always log to both memory logger (for display on the web) and
// console logger (for records / display of command line worker).
func (wi *Instance) setAndStartWorker(wrk Worker, wr *wrangler.Wrangler) (chan struct{}, error) {
	wi.currentWorkerMutex.Lock()
	defer wi.currentWorkerMutex.Unlock()
	if wi.currentWorker != nil {
		return nil, fmt.Errorf("A worker is already in progress: %v", wi.currentWorker)
	}

	wi.currentWorker = wrk
	wi.currentMemoryLogger = logutil.NewMemoryLogger()
	wi.currentContext, wi.currentCancelFunc = context.WithCancel(wi.backgroundContext)
	wi.lastRunError = nil
	done := make(chan struct{})
	wranglerLogger := wr.Logger()
	if wr == wi.wr {
		// If it's the default wrangler, do not reuse its logger because it may have been set before.
		// Resuing it would result into an endless recursion.
		wranglerLogger = logutil.NewConsoleLogger()
	}
	wr.SetLogger(logutil.NewTeeLogger(wi.currentMemoryLogger, wranglerLogger))

	// one go function runs the worker, changes state when done
	go func() {
		log.Infof("Starting worker...")
		var err error

		// Catch all panics and always save the execution state at the end.
		defer func() {
			// The recovery code is a copy of servenv.HandlePanic().
			if x := recover(); x != nil {
				err = fmt.Errorf("uncaught %v panic: %v", "vtworker", x)
			}

			wi.currentWorkerMutex.Lock()
			wi.currentContext = nil
			wi.currentCancelFunc = nil
			wi.lastRunError = err
			wi.currentWorkerMutex.Unlock()
			close(done)
		}()

		// run will take a long time
		err = wrk.Run(wi.currentContext)
	}()

	return done, nil
}

// InstallSignalHandlers installs signal handler which exit vtworker gracefully.
func (wi *Instance) InstallSignalHandlers() {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT)
	go func() {
		s := <-sigChan
		// we got a signal, notify our modules
		wi.currentWorkerMutex.Lock()
		defer wi.currentWorkerMutex.Unlock()
		if wi.currentCancelFunc != nil {
			wi.currentCancelFunc()
		} else {
			fmt.Printf("Shutting down idle worker after receiving signal: %v", s)
			os.Exit(0)
		}
	}()
}

// Reset resets the state of a finished worker. It returns an error if the worker is still running.
func (wi *Instance) Reset() error {
	wi.currentWorkerMutex.Lock()
	defer wi.currentWorkerMutex.Unlock()

	if wi.currentWorker == nil {
		return nil
	}

	// check the worker is really done
	if wi.currentContext == nil {
		wi.currentWorker = nil
		wi.currentMemoryLogger = nil
		return nil
	}

	return errors.New("worker still executing")
}
<commit_msg>vtworker: Replace fmt.Printf() with logging call.<commit_after>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package worker

import (
	"errors"
	"fmt"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/vt/logutil"
	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/wrangler"
	"golang.org/x/net/context"
)

// Instance encapsulate the execution state of vtworker.
type Instance struct {
	// Default wrangler for all operations.
	// Users can specify their own in RunCommand() e.g. the gRPC server does this.
	wr *wrangler.Wrangler

	// mutex is protecting all the following variables
	// 3 states here:
	// - no job ever ran (or reset was run): currentWorker is nil,
	// currentContext/currentCancelFunc is nil, lastRunError is nil
	// - one worker running: currentWorker is set,
	//   currentContext/currentCancelFunc is set, lastRunError is nil
	// - (at least) one worker already ran, none is running atm:
	//   currentWorker is set, currentContext is nil, lastRunError
	//   has the error returned by the worker.
	currentWorkerMutex  sync.Mutex
	currentWorker       Worker
	currentMemoryLogger *logutil.MemoryLogger
	currentContext      context.Context
	currentCancelFunc   context.CancelFunc
	lastRunError        error

	// backgroundContext is context.Background() from main() which has to be plumbed through.
	backgroundContext      context.Context
	topoServer             topo.Server
	cell                   string
	commandDisplayInterval time.Duration
}

// NewInstance creates a new Instance.
func NewInstance(ctx context.Context, ts topo.Server, cell string, commandDisplayInterval time.Duration) *Instance {
	wi := &Instance{backgroundContext: ctx, topoServer: ts, cell: cell, commandDisplayInterval: commandDisplayInterval}
	// Note: setAndStartWorker() also adds a MemoryLogger for the webserver.
	wi.wr = wi.CreateWrangler(logutil.NewConsoleLogger())
	return wi
}

// CreateWrangler creates a new wrangler using the instance specific configuration.
func (wi *Instance) CreateWrangler(logger logutil.Logger) *wrangler.Wrangler {
	return wrangler.New(logger, wi.topoServer, tmclient.NewTabletManagerClient())
}

// setAndStartWorker will set the current worker.
// We always log to both memory logger (for display on the web) and
// console logger (for records / display of command line worker).
func (wi *Instance) setAndStartWorker(wrk Worker, wr *wrangler.Wrangler) (chan struct{}, error) {
	wi.currentWorkerMutex.Lock()
	defer wi.currentWorkerMutex.Unlock()
	if wi.currentWorker != nil {
		return nil, fmt.Errorf("A worker is already in progress: %v", wi.currentWorker)
	}

	wi.currentWorker = wrk
	wi.currentMemoryLogger = logutil.NewMemoryLogger()
	wi.currentContext, wi.currentCancelFunc = context.WithCancel(wi.backgroundContext)
	wi.lastRunError = nil
	done := make(chan struct{})
	wranglerLogger := wr.Logger()
	if wr == wi.wr {
		// If it's the default wrangler, do not reuse its logger because it may have been set before.
		// Resuing it would result into an endless recursion.
		wranglerLogger = logutil.NewConsoleLogger()
	}
	wr.SetLogger(logutil.NewTeeLogger(wi.currentMemoryLogger, wranglerLogger))

	// one go function runs the worker, changes state when done
	go func() {
		log.Infof("Starting worker...")
		var err error

		// Catch all panics and always save the execution state at the end.
		defer func() {
			// The recovery code is a copy of servenv.HandlePanic().
			if x := recover(); x != nil {
				err = fmt.Errorf("uncaught %v panic: %v", "vtworker", x)
			}

			wi.currentWorkerMutex.Lock()
			wi.currentContext = nil
			wi.currentCancelFunc = nil
			wi.lastRunError = err
			wi.currentWorkerMutex.Unlock()
			close(done)
		}()

		// run will take a long time
		err = wrk.Run(wi.currentContext)
	}()

	return done, nil
}

// InstallSignalHandlers installs signal handler which exit vtworker gracefully.
func (wi *Instance) InstallSignalHandlers() {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT)
	go func() {
		s := <-sigChan
		// we got a signal, notify our modules
		wi.currentWorkerMutex.Lock()
		defer wi.currentWorkerMutex.Unlock()
		if wi.currentCancelFunc != nil {
			wi.currentCancelFunc()
		} else {
			log.Infof("Shutting down idle worker after receiving signal: %v", s)
			os.Exit(0)
		}
	}()
}

// Reset resets the state of a finished worker. It returns an error if the worker is still running.
func (wi *Instance) Reset() error {
	wi.currentWorkerMutex.Lock()
	defer wi.currentWorkerMutex.Unlock()

	if wi.currentWorker == nil {
		return nil
	}

	// check the worker is really done
	if wi.currentContext == nil {
		wi.currentWorker = nil
		wi.currentMemoryLogger = nil
		return nil
	}

	return errors.New("worker still executing")
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"math/rand"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudtrail"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSCloudTrail_basic(t *testing.T) {
	var trail cloudtrail.Trail

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSCloudTrailDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSCloudTrailConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudTrailExists("aws_cloudtrail.foobar", &trail),
					resource.TestCheckResourceAttr("aws_cloudtrail.foobar", "include_global_service_events", "true"),
				),
			},
			resource.TestStep{
				Config: testAccAWSCloudTrailConfigModified,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudTrailExists("aws_cloudtrail.foobar", &trail),
					resource.TestCheckResourceAttr("aws_cloudtrail.foobar", "s3_key_prefix", "/prefix"),
					resource.TestCheckResourceAttr("aws_cloudtrail.foobar", "include_global_service_events", "false"),
				),
			},
		},
	})
}

func testAccCheckCloudTrailExists(n string, trail *cloudtrail.Trail) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		conn := testAccProvider.Meta().(*AWSClient).cloudtrailconn
		params := cloudtrail.DescribeTrailsInput{
			TrailNameList: []*string{aws.String(rs.Primary.ID)},
		}
		resp, err := conn.DescribeTrails(&params)
		if err != nil {
			return err
		}
		if len(resp.TrailList) == 0 {
			return fmt.Errorf("Trail not found")
		}
		*trail = *resp.TrailList[0]

		return nil
	}
}

func testAccCheckAWSCloudTrailDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).cloudtrailconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_cloudtrail" {
			continue
		}

		params := cloudtrail.DescribeTrailsInput{
			TrailNameList: []*string{aws.String(rs.Primary.ID)},
		}

		resp, err := conn.DescribeTrails(&params)

		if err == nil {
			if len(resp.TrailList) != 0 &&
				*resp.TrailList[0].Name == rs.Primary.ID {
				return fmt.Errorf("CloudTrail still exists: %s", rs.Primary.ID)
			}
		}
	}

	return nil
}

var cloudTrailRandInt = rand.New(rand.NewSource(time.Now().UnixNano())).Int()

var testAccAWSCloudTrailConfig = fmt.Sprintf(`
resource "aws_cloudtrail" "foobar" {
    name = "tf-trail-foobar"
    s3_bucket_name = "${aws_s3_bucket.foo.id}"
}

resource "aws_s3_bucket" "foo" {
	bucket = "tf-test-trail-%d"
	force_destroy = true
	policy = <<POLICY
{
	"Version": "2012-10-17",
	"Statement": [
		{
			"Sid": "AWSCloudTrailAclCheck",
			"Effect": "Allow",
			"Principal": "*",
			"Action": "s3:GetBucketAcl",
			"Resource": "arn:aws:s3:::tf-test-trail-%d"
		},
		{
			"Sid": "AWSCloudTrailWrite",
			"Effect": "Allow",
			"Principal": "*",
			"Action": "s3:PutObject",
			"Resource": "arn:aws:s3:::tf-test-trail-%d/*",
			"Condition": {
				"StringEquals": {
					"s3:x-amz-acl": "bucket-owner-full-control"
				}
			}
		}
	]
}
POLICY
}
`, cloudTrailRandInt, cloudTrailRandInt, cloudTrailRandInt)

var testAccAWSCloudTrailConfigModified = fmt.Sprintf(`
resource "aws_cloudtrail" "foobar" {
    name = "tf-trail-foobar"
    s3_bucket_name = "${aws_s3_bucket.foo.id}"
    s3_key_prefix = "/prefix"
    include_global_service_events = false
}

resource "aws_s3_bucket" "foo" {
	bucket = "tf-test-trail-%d"
	force_destroy = true
	policy = <<POLICY
{
	"Version": "2012-10-17",
	"Statement": [
		{
			"Sid": "AWSCloudTrailAclCheck",
			"Effect": "Allow",
			"Principal": "*",
			"Action": "s3:GetBucketAcl",
			"Resource": "arn:aws:s3:::tf-test-trail-%d"
		},
		{
			"Sid": "AWSCloudTrailWrite",
			"Effect": "Allow",
			"Principal": "*",
			"Action": "s3:PutObject",
			"Resource": "arn:aws:s3:::tf-test-trail-%d/*",
			"Condition": {
				"StringEquals": {
					"s3:x-amz-acl": "bucket-owner-full-control"
				}
			}
		}
	]
}
POLICY
}
`, cloudTrailRandInt, cloudTrailRandInt, cloudTrailRandInt)
<commit_msg>Tests and docs for AWS CloudTrail "enable_logging"<commit_after>package aws

import (
	"fmt"
	"math/rand"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudtrail"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSCloudTrail_basic(t *testing.T) {
	var trail cloudtrail.Trail

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSCloudTrailDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSCloudTrailConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudTrailExists("aws_cloudtrail.foobar", &trail),
					resource.TestCheckResourceAttr("aws_cloudtrail.foobar", "include_global_service_events", "true"),
				),
			},
			resource.TestStep{
				Config: testAccAWSCloudTrailConfigModified,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudTrailExists("aws_cloudtrail.foobar", &trail),
					resource.TestCheckResourceAttr("aws_cloudtrail.foobar", "s3_key_prefix", "/prefix"),
					resource.TestCheckResourceAttr("aws_cloudtrail.foobar", "include_global_service_events", "false"),
				),
			},
		},
	})
}

func TestAccAWSCloudTrail_enable_logging(t *testing.T) {
	var trail cloudtrail.Trail

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSCloudTrailDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSCloudTrailConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudTrailExists("aws_cloudtrail.foobar", &trail),
					testAccCheckCloudTrailLoggingEnabled("aws_cloudtrail.foobar", false, &trail),
				),
			},
			resource.TestStep{
				Config: testAccAWSCloudTrailConfigModified,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudTrailExists("aws_cloudtrail.foobar", &trail),
					testAccCheckCloudTrailLoggingEnabled("aws_cloudtrail.foobar", true, &trail),
				),
			},
			resource.TestStep{
				Config: testAccAWSCloudTrailConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudTrailExists("aws_cloudtrail.foobar", &trail),
					testAccCheckCloudTrailLoggingEnabled("aws_cloudtrail.foobar", false, &trail),
				),
			},
		},
	})
}

func testAccCheckCloudTrailExists(n string, trail *cloudtrail.Trail) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		conn := testAccProvider.Meta().(*AWSClient).cloudtrailconn
		params := cloudtrail.DescribeTrailsInput{
			TrailNameList: []*string{aws.String(rs.Primary.ID)},
		}
		resp, err := conn.DescribeTrails(&params)
		if err != nil {
			return err
		}
		if len(resp.TrailList) == 0 {
			return fmt.Errorf("Trail not found")
		}
		*trail = *resp.TrailList[0]

		return nil
	}
}

func testAccCheckCloudTrailLoggingEnabled(n string, desired bool, trail *cloudtrail.Trail) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		conn := testAccProvider.Meta().(*AWSClient).cloudtrailconn
		params := cloudtrail.GetTrailStatusInput{
			Name: aws.String(rs.Primary.ID),
		}
		resp, err := conn.GetTrailStatus(&params)

		if err != nil {
			return err
		}
		if *resp.IsLogging != desired {
			return fmt.Errorf("Logging status is incorrect")
		}

		return nil
	}
}

func testAccCheckAWSCloudTrailDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).cloudtrailconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_cloudtrail" {
			continue
		}

		params := cloudtrail.DescribeTrailsInput{
			TrailNameList: []*string{aws.String(rs.Primary.ID)},
		}

		resp, err := conn.DescribeTrails(&params)

		if err == nil {
			if len(resp.TrailList) != 0 &&
				*resp.TrailList[0].Name == rs.Primary.ID {
				return fmt.Errorf("CloudTrail still exists: %s", rs.Primary.ID)
			}
		}
	}

	return nil
}

var cloudTrailRandInt = rand.New(rand.NewSource(time.Now().UnixNano())).Int()

var testAccAWSCloudTrailConfig = fmt.Sprintf(`
resource "aws_cloudtrail" "foobar" {
    name = "tf-trail-foobar"
    s3_bucket_name = "${aws_s3_bucket.foo.id}"
}

resource "aws_s3_bucket" "foo" {
	bucket = "tf-test-trail-%d"
	force_destroy = true
	policy = <<POLICY
{
	"Version": "2012-10-17",
	"Statement": [
		{
			"Sid": "AWSCloudTrailAclCheck",
			"Effect": "Allow",
			"Principal": "*",
			"Action": "s3:GetBucketAcl",
			"Resource": "arn:aws:s3:::tf-test-trail-%d"
		},
		{
			"Sid": "AWSCloudTrailWrite",
			"Effect": "Allow",
			"Principal": "*",
			"Action": "s3:PutObject",
			"Resource": "arn:aws:s3:::tf-test-trail-%d/*",
			"Condition": {
				"StringEquals": {
					"s3:x-amz-acl": "bucket-owner-full-control"
				}
			}
		}
	]
}
POLICY
}
`, cloudTrailRandInt, cloudTrailRandInt, cloudTrailRandInt)

var testAccAWSCloudTrailConfigModified = fmt.Sprintf(`
resource "aws_cloudtrail" "foobar" {
    name = "tf-trail-foobar"
    s3_bucket_name = "${aws_s3_bucket.foo.id}"
    s3_key_prefix = "/prefix"
    include_global_service_events = false
    enable_logging = true
}

resource "aws_s3_bucket" "foo" {
	bucket = "tf-test-trail-%d"
	force_destroy = true
	policy = <<POLICY
{
	"Version": "2012-10-17",
	"Statement": [
		{
			"Sid": "AWSCloudTrailAclCheck",
			"Effect": "Allow",
			"Principal": "*",
			"Action": "s3:GetBucketAcl",
			"Resource": "arn:aws:s3:::tf-test-trail-%d"
		},
		{
			"Sid": "AWSCloudTrailWrite",
			"Effect": "Allow",
			"Principal": "*",
			"Action": "s3:PutObject",
			"Resource": "arn:aws:s3:::tf-test-trail-%d/*",
			"Condition": {
				"StringEquals": {
					"s3:x-amz-acl": "bucket-owner-full-control"
				}
			}
		}
	]
}
POLICY
}
`, cloudTrailRandInt, cloudTrailRandInt, cloudTrailRandInt)
<|endoftext|>


<commit_before>/*
Copyright 2014 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 patch

import (
	"fmt"
	"io/ioutil"
	"reflect"
	"strings"

	jsonpatch "github.com/evanphx/json-patch"
	"github.com/spf13/cobra"
	"k8s.io/klog/v2"

	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apimachinery/pkg/util/strategicpatch"
	"k8s.io/apimachinery/pkg/util/yaml"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/cli-runtime/pkg/printers"
	"k8s.io/cli-runtime/pkg/resource"
	cmdutil "k8s.io/kubectl/pkg/cmd/util"
	"k8s.io/kubectl/pkg/scheme"
	"k8s.io/kubectl/pkg/util/i18n"
	"k8s.io/kubectl/pkg/util/templates"
)

var patchTypes = map[string]types.PatchType{"json": types.JSONPatchType, "merge": types.MergePatchType, "strategic": types.StrategicMergePatchType}

// PatchOptions is the start of the data required to perform the operation.  As new fields are added, add them here instead of
// referencing the cmd.Flags()
type PatchOptions struct {
	resource.FilenameOptions

	RecordFlags *genericclioptions.RecordFlags
	PrintFlags  *genericclioptions.PrintFlags
	ToPrinter   func(string) (printers.ResourcePrinter, error)
	Recorder    genericclioptions.Recorder

	Local     bool
	PatchType string
	Patch     string
	PatchFile string

	namespace                    string
	enforceNamespace             bool
	dryRunStrategy               cmdutil.DryRunStrategy
	dryRunVerifier               *resource.DryRunVerifier
	outputFormat                 string
	args                         []string
	builder                      *resource.Builder
	unstructuredClientForMapping func(mapping *meta.RESTMapping) (resource.RESTClient, error)
	fieldManager                 string

	genericclioptions.IOStreams
}

var (
	patchLong = templates.LongDesc(i18n.T(`
		Update field(s) of a resource using strategic merge patch, a JSON merge patch, or a JSON patch.

		JSON and YAML formats are accepted.`))

	patchExample = templates.Examples(i18n.T(`
		# Partially update a node using a strategic merge patch. Specify the patch as JSON.
		kubectl patch node k8s-node-1 -p '{"spec":{"unschedulable":true}}'

		# Partially update a node using a strategic merge patch. Specify the patch as YAML.
		kubectl patch node k8s-node-1 -p $'spec:\n unschedulable: true'

		# Partially update a node identified by the type and name specified in "node.json" using strategic merge patch.
		kubectl patch -f node.json -p '{"spec":{"unschedulable":true}}'

		# Update a container's image; spec.containers[*].name is required because it's a merge key.
		kubectl patch pod valid-pod -p '{"spec":{"containers":[{"name":"kubernetes-serve-hostname","image":"new image"}]}}'

		# Update a container's image using a json patch with positional arrays.
		kubectl patch pod valid-pod --type='json' -p='[{"op": "replace", "path": "/spec/containers/0/image", "value":"new image"}]'`))
)

func NewPatchOptions(ioStreams genericclioptions.IOStreams) *PatchOptions {
	return &PatchOptions{
		RecordFlags: genericclioptions.NewRecordFlags(),
		Recorder:    genericclioptions.NoopRecorder{},
		PrintFlags:  genericclioptions.NewPrintFlags("patched").WithTypeSetter(scheme.Scheme),
		IOStreams:   ioStreams,
	}
}

func NewCmdPatch(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	o := NewPatchOptions(ioStreams)

	cmd := &cobra.Command{
		Use:                   "patch (-f FILENAME | TYPE NAME) [-p PATCH|--patch-file FILE]",
		DisableFlagsInUseLine: true,
		Short:                 i18n.T("Update field(s) of a resource"),
		Long:                  patchLong,
		Example:               patchExample,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Complete(f, cmd, args))
			cmdutil.CheckErr(o.Validate())
			cmdutil.CheckErr(o.RunPatch())
		},
	}

	o.RecordFlags.AddFlags(cmd)
	o.PrintFlags.AddFlags(cmd)

	cmd.Flags().StringVarP(&o.Patch, "patch", "p", "", "The patch to be applied to the resource JSON file.")
	cmd.Flags().StringVar(&o.PatchFile, "patch-file", "", "A file containing a patch to be applied to the resource.")
	cmd.Flags().StringVar(&o.PatchType, "type", "strategic", fmt.Sprintf("The type of patch being provided; one of %v", sets.StringKeySet(patchTypes).List()))
	cmdutil.AddDryRunFlag(cmd)
	cmdutil.AddFilenameOptionFlags(cmd, &o.FilenameOptions, "identifying the resource to update")
	cmd.Flags().BoolVar(&o.Local, "local", o.Local, "If true, patch will operate on the content of the file, not the server-side resource.")
	cmdutil.AddFieldManagerFlagVar(cmd, &o.fieldManager, "kubectl-patch")

	return cmd
}

func (o *PatchOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
	var err error
	o.RecordFlags.Complete(cmd)
	o.Recorder, err = o.RecordFlags.ToRecorder()
	if err != nil {
		return err
	}

	o.outputFormat = cmdutil.GetFlagString(cmd, "output")
	o.dryRunStrategy, err = cmdutil.GetDryRunStrategy(cmd)
	if err != nil {
		return err
	}

	cmdutil.PrintFlagsWithDryRunStrategy(o.PrintFlags, o.dryRunStrategy)
	o.ToPrinter = func(operation string) (printers.ResourcePrinter, error) {
		o.PrintFlags.NamePrintFlags.Operation = operation

		return o.PrintFlags.ToPrinter()
	}

	o.namespace, o.enforceNamespace, err = f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}
	o.args = args
	o.builder = f.NewBuilder()
	o.unstructuredClientForMapping = f.UnstructuredClientForMapping
	dynamicClient, err := f.DynamicClient()
	if err != nil {
		return err
	}
	discoveryClient, err := f.ToDiscoveryClient()
	if err != nil {
		return err
	}
	o.dryRunVerifier = resource.NewDryRunVerifier(dynamicClient, discoveryClient)

	return nil
}

func (o *PatchOptions) Validate() error {
	if len(o.Patch) > 0 && len(o.PatchFile) > 0 {
		return fmt.Errorf("cannot specify --patch and --patch-file together")
	}
	if len(o.Patch) == 0 && len(o.PatchFile) == 0 {
		return fmt.Errorf("must specify --patch or --patch-file containing the contents of the patch")
	}
	if o.Local && len(o.args) != 0 {
		return fmt.Errorf("cannot specify --local and server resources")
	}
	if o.Local && o.dryRunStrategy == cmdutil.DryRunServer {
		return fmt.Errorf("cannot specify --local and --dry-run=server - did you mean --dry-run=client?")
	}
	if len(o.PatchType) != 0 {
		if _, ok := patchTypes[strings.ToLower(o.PatchType)]; !ok {
			return fmt.Errorf("--type must be one of %v, not %q", sets.StringKeySet(patchTypes).List(), o.PatchType)
		}
	}

	return nil
}

func (o *PatchOptions) RunPatch() error {
	patchType := types.StrategicMergePatchType
	if len(o.PatchType) != 0 {
		patchType = patchTypes[strings.ToLower(o.PatchType)]
	}

	var patchBytes []byte
	if len(o.PatchFile) > 0 {
		var err error
		patchBytes, err = ioutil.ReadFile(o.PatchFile)
		if err != nil {
			return fmt.Errorf("unable to read patch file: %v", err)
		}
	} else {
		patchBytes = []byte(o.Patch)
	}

	patchBytes, err := yaml.ToJSON(patchBytes)
	if err != nil {
		return fmt.Errorf("unable to parse %q: %v", o.Patch, err)
	}

	r := o.builder.
		Unstructured().
		ContinueOnError().
		LocalParam(o.Local).
		NamespaceParam(o.namespace).DefaultNamespace().
		FilenameParam(o.enforceNamespace, &o.FilenameOptions).
		ResourceTypeOrNameArgs(false, o.args...).
		Flatten().
		Do()
	err = r.Err()
	if err != nil {
		return err
	}

	count := 0
	err = r.Visit(func(info *resource.Info, err error) error {
		if err != nil {
			return err
		}
		count++
		name, namespace := info.Name, info.Namespace

		if !o.Local && o.dryRunStrategy != cmdutil.DryRunClient {
			mapping := info.ResourceMapping()
			if o.dryRunStrategy == cmdutil.DryRunServer {
				if err := o.dryRunVerifier.HasSupport(mapping.GroupVersionKind); err != nil {
					return err
				}
			}
			client, err := o.unstructuredClientForMapping(mapping)
			if err != nil {
				return err
			}

			helper := resource.
				NewHelper(client, mapping).
				DryRun(o.dryRunStrategy == cmdutil.DryRunServer).
				WithFieldManager(o.fieldManager)
			patchedObj, err := helper.Patch(namespace, name, patchType, patchBytes, nil)
			if err != nil {
				return err
			}

			didPatch := !reflect.DeepEqual(info.Object, patchedObj)

			// if the recorder makes a change, compute and create another patch
			if mergePatch, err := o.Recorder.MakeRecordMergePatch(patchedObj); err != nil {
				klog.V(4).Infof("error recording current command: %v", err)
			} else if len(mergePatch) > 0 {
				if recordedObj, err := helper.Patch(info.Namespace, info.Name, types.MergePatchType, mergePatch, nil); err != nil {
					klog.V(4).Infof("error recording reason: %v", err)
				} else {
					patchedObj = recordedObj
				}
			}

			printer, err := o.ToPrinter(patchOperation(didPatch))
			if err != nil {
				return err
			}
			return printer.PrintObj(patchedObj, o.Out)
		}

		originalObjJS, err := runtime.Encode(unstructured.UnstructuredJSONScheme, info.Object)
		if err != nil {
			return err
		}

		originalPatchedObjJS, err := getPatchedJSON(patchType, originalObjJS, patchBytes, info.Object.GetObjectKind().GroupVersionKind(), scheme.Scheme)
		if err != nil {
			return err
		}

		targetObj, err := runtime.Decode(unstructured.UnstructuredJSONScheme, originalPatchedObjJS)
		if err != nil {
			return err
		}

		didPatch := !reflect.DeepEqual(info.Object, targetObj)
		printer, err := o.ToPrinter(patchOperation(didPatch))
		if err != nil {
			return err
		}
		return printer.PrintObj(targetObj, o.Out)
	})
	if err != nil {
		return err
	}
	if count == 0 {
		return fmt.Errorf("no objects passed to patch")
	}
	return nil
}

func getPatchedJSON(patchType types.PatchType, originalJS, patchJS []byte, gvk schema.GroupVersionKind, creater runtime.ObjectCreater) ([]byte, error) {
	switch patchType {
	case types.JSONPatchType:
		patchObj, err := jsonpatch.DecodePatch(patchJS)
		if err != nil {
			return nil, err
		}
		bytes, err := patchObj.Apply(originalJS)
		// TODO: This is pretty hacky, we need a better structured error from the json-patch
		if err != nil && strings.Contains(err.Error(), "doc is missing key") {
			msg := err.Error()
			ix := strings.Index(msg, "key:")
			key := msg[ix+5:]
			return bytes, fmt.Errorf("Object to be patched is missing field (%s)", key)
		}
		return bytes, err

	case types.MergePatchType:
		return jsonpatch.MergePatch(originalJS, patchJS)

	case types.StrategicMergePatchType:
		// get a typed object for this GVK if we need to apply a strategic merge patch
		obj, err := creater.New(gvk)
		if err != nil {
			return nil, fmt.Errorf("cannot apply strategic merge patch for %s locally, try --type merge", gvk.String())
		}
		return strategicpatch.StrategicMergePatch(originalJS, patchJS, obj)

	default:
		// only here as a safety net - go-restful filters content-type
		return nil, fmt.Errorf("unknown Content-Type header for patch: %v", patchType)
	}
}

func patchOperation(didPatch bool) string {
	if didPatch {
		return "patched"
	}
	return "patched (no change)"
}
<commit_msg>kubectl: fix error message for client-side SMP for CRs<commit_after>/*
Copyright 2014 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 patch

import (
	"fmt"
	"io/ioutil"
	"reflect"
	"strings"

	jsonpatch "github.com/evanphx/json-patch"
	"github.com/spf13/cobra"
	"k8s.io/klog/v2"

	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apimachinery/pkg/util/strategicpatch"
	"k8s.io/apimachinery/pkg/util/yaml"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/cli-runtime/pkg/printers"
	"k8s.io/cli-runtime/pkg/resource"
	cmdutil "k8s.io/kubectl/pkg/cmd/util"
	"k8s.io/kubectl/pkg/scheme"
	"k8s.io/kubectl/pkg/util/i18n"
	"k8s.io/kubectl/pkg/util/templates"
)

var patchTypes = map[string]types.PatchType{"json": types.JSONPatchType, "merge": types.MergePatchType, "strategic": types.StrategicMergePatchType}

// PatchOptions is the start of the data required to perform the operation.  As new fields are added, add them here instead of
// referencing the cmd.Flags()
type PatchOptions struct {
	resource.FilenameOptions

	RecordFlags *genericclioptions.RecordFlags
	PrintFlags  *genericclioptions.PrintFlags
	ToPrinter   func(string) (printers.ResourcePrinter, error)
	Recorder    genericclioptions.Recorder

	Local     bool
	PatchType string
	Patch     string
	PatchFile string

	namespace                    string
	enforceNamespace             bool
	dryRunStrategy               cmdutil.DryRunStrategy
	dryRunVerifier               *resource.DryRunVerifier
	outputFormat                 string
	args                         []string
	builder                      *resource.Builder
	unstructuredClientForMapping func(mapping *meta.RESTMapping) (resource.RESTClient, error)
	fieldManager                 string

	genericclioptions.IOStreams
}

var (
	patchLong = templates.LongDesc(i18n.T(`
		Update field(s) of a resource using strategic merge patch, a JSON merge patch, or a JSON patch.

		JSON and YAML formats are accepted.`))

	patchExample = templates.Examples(i18n.T(`
		# Partially update a node using a strategic merge patch. Specify the patch as JSON.
		kubectl patch node k8s-node-1 -p '{"spec":{"unschedulable":true}}'

		# Partially update a node using a strategic merge patch. Specify the patch as YAML.
		kubectl patch node k8s-node-1 -p $'spec:\n unschedulable: true'

		# Partially update a node identified by the type and name specified in "node.json" using strategic merge patch.
		kubectl patch -f node.json -p '{"spec":{"unschedulable":true}}'

		# Update a container's image; spec.containers[*].name is required because it's a merge key.
		kubectl patch pod valid-pod -p '{"spec":{"containers":[{"name":"kubernetes-serve-hostname","image":"new image"}]}}'

		# Update a container's image using a json patch with positional arrays.
		kubectl patch pod valid-pod --type='json' -p='[{"op": "replace", "path": "/spec/containers/0/image", "value":"new image"}]'`))
)

func NewPatchOptions(ioStreams genericclioptions.IOStreams) *PatchOptions {
	return &PatchOptions{
		RecordFlags: genericclioptions.NewRecordFlags(),
		Recorder:    genericclioptions.NoopRecorder{},
		PrintFlags:  genericclioptions.NewPrintFlags("patched").WithTypeSetter(scheme.Scheme),
		IOStreams:   ioStreams,
	}
}

func NewCmdPatch(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	o := NewPatchOptions(ioStreams)

	cmd := &cobra.Command{
		Use:                   "patch (-f FILENAME | TYPE NAME) [-p PATCH|--patch-file FILE]",
		DisableFlagsInUseLine: true,
		Short:                 i18n.T("Update field(s) of a resource"),
		Long:                  patchLong,
		Example:               patchExample,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Complete(f, cmd, args))
			cmdutil.CheckErr(o.Validate())
			cmdutil.CheckErr(o.RunPatch())
		},
	}

	o.RecordFlags.AddFlags(cmd)
	o.PrintFlags.AddFlags(cmd)

	cmd.Flags().StringVarP(&o.Patch, "patch", "p", "", "The patch to be applied to the resource JSON file.")
	cmd.Flags().StringVar(&o.PatchFile, "patch-file", "", "A file containing a patch to be applied to the resource.")
	cmd.Flags().StringVar(&o.PatchType, "type", "strategic", fmt.Sprintf("The type of patch being provided; one of %v", sets.StringKeySet(patchTypes).List()))
	cmdutil.AddDryRunFlag(cmd)
	cmdutil.AddFilenameOptionFlags(cmd, &o.FilenameOptions, "identifying the resource to update")
	cmd.Flags().BoolVar(&o.Local, "local", o.Local, "If true, patch will operate on the content of the file, not the server-side resource.")
	cmdutil.AddFieldManagerFlagVar(cmd, &o.fieldManager, "kubectl-patch")

	return cmd
}

func (o *PatchOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
	var err error
	o.RecordFlags.Complete(cmd)
	o.Recorder, err = o.RecordFlags.ToRecorder()
	if err != nil {
		return err
	}

	o.outputFormat = cmdutil.GetFlagString(cmd, "output")
	o.dryRunStrategy, err = cmdutil.GetDryRunStrategy(cmd)
	if err != nil {
		return err
	}

	cmdutil.PrintFlagsWithDryRunStrategy(o.PrintFlags, o.dryRunStrategy)
	o.ToPrinter = func(operation string) (printers.ResourcePrinter, error) {
		o.PrintFlags.NamePrintFlags.Operation = operation

		return o.PrintFlags.ToPrinter()
	}

	o.namespace, o.enforceNamespace, err = f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}
	o.args = args
	o.builder = f.NewBuilder()
	o.unstructuredClientForMapping = f.UnstructuredClientForMapping
	dynamicClient, err := f.DynamicClient()
	if err != nil {
		return err
	}
	discoveryClient, err := f.ToDiscoveryClient()
	if err != nil {
		return err
	}
	o.dryRunVerifier = resource.NewDryRunVerifier(dynamicClient, discoveryClient)

	return nil
}

func (o *PatchOptions) Validate() error {
	if len(o.Patch) > 0 && len(o.PatchFile) > 0 {
		return fmt.Errorf("cannot specify --patch and --patch-file together")
	}
	if len(o.Patch) == 0 && len(o.PatchFile) == 0 {
		return fmt.Errorf("must specify --patch or --patch-file containing the contents of the patch")
	}
	if o.Local && len(o.args) != 0 {
		return fmt.Errorf("cannot specify --local and server resources")
	}
	if o.Local && o.dryRunStrategy == cmdutil.DryRunServer {
		return fmt.Errorf("cannot specify --local and --dry-run=server - did you mean --dry-run=client?")
	}
	if len(o.PatchType) != 0 {
		if _, ok := patchTypes[strings.ToLower(o.PatchType)]; !ok {
			return fmt.Errorf("--type must be one of %v, not %q", sets.StringKeySet(patchTypes).List(), o.PatchType)
		}
	}

	return nil
}

func (o *PatchOptions) RunPatch() error {
	patchType := types.StrategicMergePatchType
	if len(o.PatchType) != 0 {
		patchType = patchTypes[strings.ToLower(o.PatchType)]
	}

	var patchBytes []byte
	if len(o.PatchFile) > 0 {
		var err error
		patchBytes, err = ioutil.ReadFile(o.PatchFile)
		if err != nil {
			return fmt.Errorf("unable to read patch file: %v", err)
		}
	} else {
		patchBytes = []byte(o.Patch)
	}

	patchBytes, err := yaml.ToJSON(patchBytes)
	if err != nil {
		return fmt.Errorf("unable to parse %q: %v", o.Patch, err)
	}

	r := o.builder.
		Unstructured().
		ContinueOnError().
		LocalParam(o.Local).
		NamespaceParam(o.namespace).DefaultNamespace().
		FilenameParam(o.enforceNamespace, &o.FilenameOptions).
		ResourceTypeOrNameArgs(false, o.args...).
		Flatten().
		Do()
	err = r.Err()
	if err != nil {
		return err
	}

	count := 0
	err = r.Visit(func(info *resource.Info, err error) error {
		if err != nil {
			return err
		}
		count++
		name, namespace := info.Name, info.Namespace

		if !o.Local && o.dryRunStrategy != cmdutil.DryRunClient {
			mapping := info.ResourceMapping()
			if o.dryRunStrategy == cmdutil.DryRunServer {
				if err := o.dryRunVerifier.HasSupport(mapping.GroupVersionKind); err != nil {
					return err
				}
			}
			client, err := o.unstructuredClientForMapping(mapping)
			if err != nil {
				return err
			}

			helper := resource.
				NewHelper(client, mapping).
				DryRun(o.dryRunStrategy == cmdutil.DryRunServer).
				WithFieldManager(o.fieldManager)
			patchedObj, err := helper.Patch(namespace, name, patchType, patchBytes, nil)
			if err != nil {
				return err
			}

			didPatch := !reflect.DeepEqual(info.Object, patchedObj)

			// if the recorder makes a change, compute and create another patch
			if mergePatch, err := o.Recorder.MakeRecordMergePatch(patchedObj); err != nil {
				klog.V(4).Infof("error recording current command: %v", err)
			} else if len(mergePatch) > 0 {
				if recordedObj, err := helper.Patch(info.Namespace, info.Name, types.MergePatchType, mergePatch, nil); err != nil {
					klog.V(4).Infof("error recording reason: %v", err)
				} else {
					patchedObj = recordedObj
				}
			}

			printer, err := o.ToPrinter(patchOperation(didPatch))
			if err != nil {
				return err
			}
			return printer.PrintObj(patchedObj, o.Out)
		}

		originalObjJS, err := runtime.Encode(unstructured.UnstructuredJSONScheme, info.Object)
		if err != nil {
			return err
		}

		originalPatchedObjJS, err := getPatchedJSON(patchType, originalObjJS, patchBytes, info.Object.GetObjectKind().GroupVersionKind(), scheme.Scheme)
		if err != nil {
			return err
		}

		targetObj, err := runtime.Decode(unstructured.UnstructuredJSONScheme, originalPatchedObjJS)
		if err != nil {
			return err
		}

		didPatch := !reflect.DeepEqual(info.Object, targetObj)
		printer, err := o.ToPrinter(patchOperation(didPatch))
		if err != nil {
			return err
		}
		return printer.PrintObj(targetObj, o.Out)
	})
	if err != nil {
		return err
	}
	if count == 0 {
		return fmt.Errorf("no objects passed to patch")
	}
	return nil
}

func getPatchedJSON(patchType types.PatchType, originalJS, patchJS []byte, gvk schema.GroupVersionKind, creater runtime.ObjectCreater) ([]byte, error) {
	switch patchType {
	case types.JSONPatchType:
		patchObj, err := jsonpatch.DecodePatch(patchJS)
		if err != nil {
			return nil, err
		}
		bytes, err := patchObj.Apply(originalJS)
		// TODO: This is pretty hacky, we need a better structured error from the json-patch
		if err != nil && strings.Contains(err.Error(), "doc is missing key") {
			msg := err.Error()
			ix := strings.Index(msg, "key:")
			key := msg[ix+5:]
			return bytes, fmt.Errorf("Object to be patched is missing field (%s)", key)
		}
		return bytes, err

	case types.MergePatchType:
		return jsonpatch.MergePatch(originalJS, patchJS)

	case types.StrategicMergePatchType:
		// get a typed object for this GVK if we need to apply a strategic merge patch
		obj, err := creater.New(gvk)
		if err != nil {
			return nil, fmt.Errorf("strategic merge patch is not supported for %s locally, try --type merge", gvk.String())
		}
		return strategicpatch.StrategicMergePatch(originalJS, patchJS, obj)

	default:
		// only here as a safety net - go-restful filters content-type
		return nil, fmt.Errorf("unknown Content-Type header for patch: %v", patchType)
	}
}

func patchOperation(didPatch bool) string {
	if didPatch {
		return "patched"
	}
	return "patched (no change)"
}
<|endoftext|>


<commit_before>// Package main runs a Cloud Function locally.
package main

import (
	"context"
	"log"

	"cloud.google.com/go/pubsub"
	exec "github.com/googlecloudplatform/security-response-automation"
)

const (
	badIPCSCC = `{
		"notificationConfigName": "organizations/1037840971520/notificationConfigs/sampleConfigId",
		"finding": {
		  "name": "organizations/1037840971520/sources/15233230630886231666/findings/6cc800d88324478aa80d2031794214a4",
		  "parent": "organizations/1037840971520/sources/15233230630886231666",
		  "resourceName": "//cloudresourcemanager.googleapis.com/projects/459837319394",
		  "state": "ACTIVE",
		  "category": "C2: Bad IP",
		  "sourceProperties": {
			"properties_subnetwork_name": "default",
			"detectionCategory_ruleName": "bad_ip",
			"properties_project_id": "ae-threat-detection",
			"eventTime": "2019-12-11T22:39:54.657Z",
			"sourceId_projectNumber": "459837319394",
			"affectedResources_0_gcpResourceName": "//cloudresourcemanager.googleapis.com/projects/459837319394",
			"properties_ip_0": "80.82.64.214",
			"properties_location": "us-central1-a",
			"evidence_0_sourceLogId_timestamp": "2019-12-11T22:39:54.060661565Z",
			"detectionPriority": "HIGH",
			"properties_instanceDetails": "/projects/ae-threat-detection/zones/us-central1-a/instances/instance-1",
			"detectionCategory_technique": "C2",
			"findingId": "6cc800d88324478aa80d2031794214a4",
			"sourceId_customerOrganizationNumber": "1037840971520",
			"evidence_0_sourceLogId_insertId": "198cht8g3wedctq",
			"detectionCategory_indicator": "ip",
			"properties_subnetwork_id": "3951118263795572377"
		  },
		  "securityMarks": {
			"name": "organizations/1037840971520/sources/15233230630886231666/findings/6cc800d88324478aa80d2031794214a4/securityMarks",
			"marks": {
			  "k": "ok",
			  "adfs": "adf",
			  "adsf": "adsf"
			}
		  },
		  "eventTime": "2019-12-11T22:39:54.657Z",
		  "createTime": "2019-12-11T22:39:55.358Z"
		}
	}`
	badIP = `{
		"jsonPayload": {
			"properties": {
				"location": "us-central1",
				"project_id": "aerial-jigsaw-235219",
				"instanceDetails": "/zones/us-central1-a/instances/instance-4"
			},
			"detectionCategory": {
				"ruleName": "bad_ip"
			}
		},
		"logName": "projects/test-project/logs/threatdetection.googleapis.com` + "%%2F" + `detection"
	}`
	badIPValues = `{
		"ProjectID": "aerial-jigsaw-235219",
		"Instance": "instance-4",
		"Zone": "us-central1-a",
		"RuleName": "bad_ip",
		"DryRun": true,
		"Target": [],
		"Exclude": [],
		"Output": []
	}`
	sshBruteForce = `{
		"jsonPayload": {
		  "properties": {
				"location": "us-central1",
				"project_id": "dark-shade",
				"instanceDetails": "/zones/us-central1-c/instances/instance-",
				"project_id": "aerial-jigsaw-235219",
				"loginAttempts": [{
						"authResult": "FAIL",
						"sourceIp": "10.200.0.2",
						"userName": "okokok",
						"vmName": "ssh-password-auth-debian-9"
					}, {
						"authResult": "SUCCESS",
						"sourceIp": "10.200.0.2",
						"userName": "okokok",
						"vmName": "ssh-password-auth-debian-9"
					}]
		  	},
		  "detectionCategory": {
			"ruleName": "ssh_brute_force"
		  }
		},
		"logName": "projects/test-project/logs/threatdetection.googleapis.com` + "%%2F" + `detection"
	}`
	iam = `{
		"jsonPayload": {
		  "affectedResources": [
			{
			  "gcpResourceName": "//cloudresourcemanager.googleapis.com/projects/997507777601"
			}
		  ],
		  "detectionCategory": {
			"indicator": "audit_log",
			"ruleName": "iam_anomalous_grant",
			"subRuleName": "external_member_added_to_policy",
			"technique": "persistence"
		  },
		  "detectionPriority": "HIGH",
		  "eventTime": "2019-12-12T22:42:36.720Z",
		  "evidence": [
			{
			  "sourceLogId": {
				"insertId": "28alkud1zr2",
				"timestamp": "2019-12-12T22:42:35.113Z"
			  }
			}
		  ],
		  "findingId": "829f7e10a7aa43f48c91c23407d3e16b",
		  "properties": {
			"bindingDeltas": [
			  {
				"action": "ADD",
				"member": "user:ccexperts@gmail.com",
				"role": "roles/editor"
			  }
			],
			"externalMembers": [
			  "user:ccexperts@gmail.com"
			],
			"principalEmail": "tom3fitzgerald@gmail.com",
			"project_id": "aerial-jigsaw-235219"
		  },
		  "sourceId": {
			"customerOrganizationNumber": "154584661726",
			"projectNumber": "997507777601"
		  }
		},
		"logName": "projects/aerial-jigsaw-235219/logs/threatdetection.googleapis.com%2Fdetection",
		"receiveTimestamp": "2019-12-12T22:42:37.613334916Z",
		"resource": {
		  "labels": {
			"detector_name": "iam_anomalous_grant",
			"project_id": "aerial-jigsaw-235219"
		  },
		  "type": "threat_detector"
		},
		"severity": "CRITICAL",
		"timestamp": "2019-12-12T22:42:36.720Z"
	  }`
)

func main() {
	ctx := context.Background()
	// if err := exec.IAMRevoke(ctx, pubsub.Message{Data: []byte(`{
	// 	"ProjectID": "aerial-jigsaw-235219",
	// 	"ExternalMembers": [
	// 		"ccexperts@gmail.com"
	// 	]
	// }`)}); err != nil {
	// 	log.Fatal(err)
	// }
	// if err := exec.SnapshotDisk(ctx, pubsub.Message{Data: []byte(badIP)}); err != nil {
	// 	log.Fatal(err)
	// }
	// if err := exec.OpenFirewall(ctx, pubsub.Message{Data: []byte(sshBruteForce)}); err != nil {
	// 	log.Fatal(err)
	// }
	if err := exec.Router(ctx, pubsub.Message{Data: []byte(iam)}); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>add publicBucket to cli<commit_after>// Package main runs a Cloud Function locally.
package main

import (
	"context"
	"log"

	"cloud.google.com/go/pubsub"
	exec "github.com/googlecloudplatform/security-response-automation"
)

const (
	badIPCSCC = `{
		"notificationConfigName": "organizations/1037840971520/notificationConfigs/sampleConfigId",
		"finding": {
		  "name": "organizations/1037840971520/sources/15233230630886231666/findings/6cc800d88324478aa80d2031794214a4",
		  "parent": "organizations/1037840971520/sources/15233230630886231666",
		  "resourceName": "//cloudresourcemanager.googleapis.com/projects/459837319394",
		  "state": "ACTIVE",
		  "category": "C2: Bad IP",
		  "sourceProperties": {
			"properties_subnetwork_name": "default",
			"detectionCategory_ruleName": "bad_ip",
			"properties_project_id": "ae-threat-detection",
			"eventTime": "2019-12-11T22:39:54.657Z",
			"sourceId_projectNumber": "459837319394",
			"affectedResources_0_gcpResourceName": "//cloudresourcemanager.googleapis.com/projects/459837319394",
			"properties_ip_0": "80.82.64.214",
			"properties_location": "us-central1-a",
			"evidence_0_sourceLogId_timestamp": "2019-12-11T22:39:54.060661565Z",
			"detectionPriority": "HIGH",
			"properties_instanceDetails": "/projects/ae-threat-detection/zones/us-central1-a/instances/instance-1",
			"detectionCategory_technique": "C2",
			"findingId": "6cc800d88324478aa80d2031794214a4",
			"sourceId_customerOrganizationNumber": "1037840971520",
			"evidence_0_sourceLogId_insertId": "198cht8g3wedctq",
			"detectionCategory_indicator": "ip",
			"properties_subnetwork_id": "3951118263795572377"
		  },
		  "securityMarks": {
			"name": "organizations/1037840971520/sources/15233230630886231666/findings/6cc800d88324478aa80d2031794214a4/securityMarks",
			"marks": {
			  "k": "ok",
			  "adfs": "adf",
			  "adsf": "adsf"
			}
		  },
		  "eventTime": "2019-12-11T22:39:54.657Z",
		  "createTime": "2019-12-11T22:39:55.358Z"
		}
	}`
	badIP = `{
		"jsonPayload": {
			"properties": {
				"location": "us-central1",
				"project_id": "aerial-jigsaw-235219",
				"instanceDetails": "/zones/us-central1-a/instances/instance-4"
			},
			"detectionCategory": {
				"ruleName": "bad_ip"
			}
		},
		"logName": "projects/test-project/logs/threatdetection.googleapis.com` + "%%2F" + `detection"
	}`
	badIPValues = `{
		"ProjectID": "aerial-jigsaw-235219",
		"Instance": "instance-4",
		"Zone": "us-central1-a",
		"RuleName": "bad_ip",
		"DryRun": true,
		"Target": [],
		"Exclude": [],
		"Output": []
	}`
	sshBruteForce = `{
		"jsonPayload": {
		  "properties": {
				"location": "us-central1",
				"project_id": "dark-shade",
				"instanceDetails": "/zones/us-central1-c/instances/instance-",
				"project_id": "aerial-jigsaw-235219",
				"loginAttempts": [{
						"authResult": "FAIL",
						"sourceIp": "10.200.0.2",
						"userName": "okokok",
						"vmName": "ssh-password-auth-debian-9"
					}, {
						"authResult": "SUCCESS",
						"sourceIp": "10.200.0.2",
						"userName": "okokok",
						"vmName": "ssh-password-auth-debian-9"
					}]
		  	},
		  "detectionCategory": {
			"ruleName": "ssh_brute_force"
		  }
		},
		"logName": "projects/test-project/logs/threatdetection.googleapis.com` + "%%2F" + `detection"
	}`
	iam = `{
		"jsonPayload": {
		  "affectedResources": [
			{
			  "gcpResourceName": "//cloudresourcemanager.googleapis.com/projects/997507777601"
			}
		  ],
		  "detectionCategory": {
			"indicator": "audit_log",
			"ruleName": "iam_anomalous_grant",
			"subRuleName": "external_member_added_to_policy",
			"technique": "persistence"
		  },
		  "detectionPriority": "HIGH",
		  "eventTime": "2019-12-12T22:42:36.720Z",
		  "evidence": [
			{
			  "sourceLogId": {
				"insertId": "28alkud1zr2",
				"timestamp": "2019-12-12T22:42:35.113Z"
			  }
			}
		  ],
		  "findingId": "829f7e10a7aa43f48c91c23407d3e16b",
		  "properties": {
			"bindingDeltas": [
			  {
				"action": "ADD",
				"member": "user:ccexperts@gmail.com",
				"role": "roles/editor"
			  }
			],
			"externalMembers": [
			  "user:ccexperts@gmail.com"
			],
			"principalEmail": "tom3fitzgerald@gmail.com",
			"project_id": "aerial-jigsaw-235219"
		  },
		  "sourceId": {
			"customerOrganizationNumber": "154584661726",
			"projectNumber": "997507777601"
		  }
		},
		"logName": "projects/aerial-jigsaw-235219/logs/threatdetection.googleapis.com%2Fdetection",
		"receiveTimestamp": "2019-12-12T22:42:37.613334916Z",
		"resource": {
		  "labels": {
			"detector_name": "iam_anomalous_grant",
			"project_id": "aerial-jigsaw-235219"
		  },
		  "type": "threat_detector"
		},
		"severity": "CRITICAL",
		"timestamp": "2019-12-12T22:42:36.720Z"
	  }`
	publicBucket = `{
		"notificationConfigName": "organizations/154584661726/notificationConfigs/sampleConfigId",
		"finding": {
			"name": "organizations/154584661726/sources/2673592633662526977/findings/782e52631d61da6117a3772137c270d8",
			"parent": "organizations/154584661726/sources/2673592633662526977",
			"resourceName": "//storage.googleapis.com/this-is-public-on-purpose",
			"state": "ACTIVE",
			"category": "PUBLIC_BUCKET_ACL",
			"externalUri": "https://console.cloud.google.com/storage/browser/this-is-public-on-purpose",
			"sourceProperties": {
				"ReactivationCount": 0.0,
				"ExceptionInstructions": "Add the security mark \"allow_public_bucket_acl\" to the asset with a value of \"true\" to prevent this finding from being activated again.",
				"SeverityLevel": "High",
				"Recommendation": "Go to https://console.cloud.google.com/storage/browser/this-is-public-on-purpose, click on the Permissions tab, and remove \"allUsers\" and \"allAuthenticatedUsers\" from the bucket's members.",
				"ProjectId": "aerial-jigsaw-235219",
				"AssetCreationTime": "2019-09-19T20:08:29.102Z",
				"ScannerName": "STORAGE_SCANNER",
				"ScanRunId": "2019-09-23T10:20:27.204-07:00",
				"Explanation": "This bucket is public and can be accessed by anyone on the Internet. \"allUsers\" represents anyone on the Internet, and \"allAuthenticatedUsers\" represents anyone who is authenticated with a Google account; neither is constrained to users within your organization."
			},
			"securityMarks": {
				"name": "organizations/154584661726/sources/2673592633662526977/findings/782e52631d61da6117a3772137c270d8/securityMarks",
				"marks": {
					"babab": "3"
				}
			},
			"eventTime": "2019-09-23T17:20:27.204Z",
			"createTime": "2019-09-23T17:20:27.934Z"
		}
	}`
)

func main() {
	ctx := context.Background()
	// if err := exec.IAMRevoke(ctx, pubsub.Message{Data: []byte(`{
	// 	"ProjectID": "aerial-jigsaw-235219",
	// 	"ExternalMembers": [
	// 		"ccexperts@gmail.com"
	// 	]
	// }`)}); err != nil {
	// 	log.Fatal(err)
	// }
	// if err := exec.SnapshotDisk(ctx, pubsub.Message{Data: []byte(badIP)}); err != nil {
	// 	log.Fatal(err)
	// }
	// if err := exec.OpenFirewall(ctx, pubsub.Message{Data: []byte(sshBruteForce)}); err != nil {
	// 	log.Fatal(err)
	// }
	if err := exec.Router(ctx, pubsub.Message{Data: []byte(publicBucket)}); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2012 Tetsuo Kiso. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package pegasos

// an simple implementation of feature vector in go

type Node struct {
	id int
	v  float64
}

func (f Node) Id() int        { return f.id }
func (f Node) Value() float64 { return f.v }

type FeatureVector struct {
	vec []Node
	ptr int // index of new element to be inserted
}

func (fv *FeatureVector) Len() int  { return len(fv.vec) }
func (fv *FeatureVector) Cap() int  { return cap(fv.vec) }
func (fv *FeatureVector) Size() int { return fv.ptr }

func (fv *FeatureVector) Index(i int) *Node {
	if i < 0 {
		panic("invalid id")
	}
	if i > fv.Len() {
		panic("index exceeds the size of feature vector")
	}
	return &fv.vec[i]
}

// Alloc
func NewFeatureVector(size int) *FeatureVector {
	v := make([]Node, size)
	fv := FeatureVector{vec: v}
	return &fv
}

func (fv *FeatureVector) PushBack(f Node) {
	l := fv.Len()

	if l > fv.Cap() { // reallocate
		// Allocate double what's needed
		newSlice := make([]Node, l*2)

		copy(newSlice, fv.vec)
		fv.vec = newSlice
	}
	fv.vec[fv.ptr] = f
	fv.ptr++
}

// Append
// func (fv *FeatureVector) Append(slice []Feature) {
// TODO: implement append slice of Feature to fv.vec.
// }

// TODO: this is too slow. profiling is needed.
func InnerProduct(w []float64, fv *FeatureVector) float64 {
	res := 0.0
	N := fv.Size()
	for i := 0; i < N; i++ {
		f := fv.Index(i)
		res += w[f.id] * f.v
	}
	return res
}
<commit_msg>Fix reallocation of feature vectors.<commit_after>// Copyright 2012 Tetsuo Kiso. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package pegasos

// an simple implementation of feature vector in go

type Node struct {
	id int
	v  float64
}

func (f Node) Id() int        { return f.id }
func (f Node) Value() float64 { return f.v }

type FeatureVector struct {
	vec []Node
	ptr int // index of new element to be inserted
}

func (fv *FeatureVector) Len() int  { return len(fv.vec) }
func (fv *FeatureVector) Cap() int  { return cap(fv.vec) }
func (fv *FeatureVector) Size() int { return fv.ptr }

func (fv *FeatureVector) Index(i int) *Node {
	if i < 0 {
		panic("invalid id")
	}
	if i > fv.Len() {
		panic("index exceeds the size of feature vector")
	}
	return &fv.vec[i]
}

// Alloc
func NewFeatureVector(size int) *FeatureVector {
	v := make([]Node, size)
	fv := FeatureVector{vec: v}
	return &fv
}

func (fv *FeatureVector) PushBack(f Node) {
	if fv.ptr >= fv.Cap() { // reallocate
		// Allocate double what's needed
		newSlice := make([]Node, fv.ptr*2)

		copy(newSlice, fv.vec)
		fv.vec = newSlice
	}
	fv.vec[fv.ptr] = f
	fv.ptr++
}

// Append
// func (fv *FeatureVector) Append(slice []Feature) {
// TODO: implement append slice of Feature to fv.vec.
// }

// TODO: this is too slow. profiling is needed.
func InnerProduct(w []float64, fv *FeatureVector) float64 {
	res := 0.0
	N := fv.Size()
	for i := 0; i < N; i++ {
		f := fv.Index(i)
		res += w[f.id] * f.v
	}
	return res
}
<|endoftext|>


<commit_before>package util

import (
  "github.com/twitchyliquid64/CNC/logging"
  "github.com/robertkrimen/otto"
  "fmt"
)

func GetFunc(argument otto.Value, vm *otto.Otto) otto.Value {
  if (argument.IsFunction()) {
    return argument;
  }

  mname := argument.String()
  logging.Warning("Using a string to name parameters is now depricated. Please use callbacks directly: " + mname)

  method, err := vm.Get(mname)
  if (err != nil){
    logging.Error("Error getting function \"" + mname + "\": ")
  }

  return method;
}
<commit_msg>Removed unused import<commit_after>package util

import (
  "github.com/twitchyliquid64/CNC/logging"
  "github.com/robertkrimen/otto"
)

func GetFunc(argument otto.Value, vm *otto.Otto) otto.Value {
  if (argument.IsFunction()) {
    return argument;
  }

  mname := argument.String()
  logging.Warning("Using a string to name parameters is now depricated. Please use callbacks directly: " + mname)

  method, err := vm.Get(mname)
  if (err != nil){
    logging.Error("Error getting function \"" + mname + "\": ")
  }

  return method;
}
<|endoftext|>