func main() { log.Infof("starting schema...") options := struct { Driver string `goptions:"-t,--type, obligatory, description='Type of database backend'"` DSN string `goptions:"-d,--database, obligatory, description='DSN of the database backend'"` }{ // No defaults } goptions.ParseAndFail(&options) database := &db.DB{ Driver: options.Driver, DSN: options.DSN, } log.Debugf("connecting to %s database at %s", database.Driver, database.DSN) if err := database.Connect(); err != nil { log.Errorf("failed to connect to %s database at %s: %s", database.Driver, database.DSN, err) } if err := database.Setup(); err != nil { log.Errorf("failed to set up schema in %s database at %s: %s", database.Driver, database.DSN, err) return } log.Infof("deployed schema version %d", db.CurrentSchema) }
func (s *Supervisor) PurgeArchives() { log.Debugf("scanning for archives needing to be expired") // mark archives past their retention policy as expired toExpire, err := s.Database.GetExpiredArchives() if err != nil { log.Errorf("error retrieving archives needing to be expired: %s", err.Error()) } for _, archive := range toExpire { log.Infof("marking archive %s has expiration date %s, marking as expired", archive.UUID, archive.ExpiresAt) err := s.Database.ExpireArchive(archive.UUID) if err != nil { log.Errorf("error marking archive %s as expired: %s", archive.UUID, err) continue } } // get archives that are not valid or purged toPurge, err := s.Database.GetArchivesNeedingPurge() if err != nil { log.Errorf("error retrieving archives to purge: %s", err.Error()) } for _, archive := range toPurge { log.Infof("requesting purge of archive %s due to status '%s'", archive.UUID, archive.Status) task, err := s.Database.CreatePurgeTask("system", archive, s.PurgeAgent) if err != nil { log.Errorf("error scheduling purge of archive %s: %s", archive.UUID, err) continue } s.ScheduleTask(task) } }
func main() { log.Infof("starting schema...") options := struct { Help bool `goptions:"-h, --help, description='Show the help screen'"` Driver string `goptions:"-t, --type, description='Type of database backend (postgres or mysql)'"` DSN string `goptions:"-d,--database, description='DSN of the database backend'"` Version bool `goptions:"-v, --version, description='Display the SHIELD version'"` }{ // No defaults } if err := goptions.Parse(&options); err != nil { fmt.Printf("%s\n", err) goptions.PrintHelp() return } if options.Help { goptions.PrintHelp() os.Exit(0) } if options.Version { if Version == "" { fmt.Printf("shield-schema (development)%s\n", Version) } else { fmt.Printf("shield-schema v%s\n", Version) } os.Exit(0) } if options.Driver == "" { fmt.Fprintf(os.Stderr, "You must indicate which type of database you wish to manage, via the `--type` option.\n") os.Exit(1) } if options.DSN == "" { fmt.Fprintf(os.Stderr, "You must specify the DSN of your database, via the `--database` option.\n") os.Exit(1) } database := &db.DB{ Driver: options.Driver, DSN: options.DSN, } log.Debugf("connecting to %s database at %s", database.Driver, database.DSN) if err := database.Connect(); err != nil { log.Errorf("failed to connect to %s database at %s: %s", database.Driver, database.DSN, err) } if err := database.Setup(); err != nil { log.Errorf("failed to set up schema in %s database at %s: %s", database.Driver, database.DSN, err) return } log.Infof("deployed schema version %d", db.CurrentSchema) }
func main() { var opts ShieldAgentOpts opts.Log = "Info" if err := goptions.Parse(&opts); err != nil { fmt.Printf("%s\n", err) goptions.PrintHelp() return } if opts.Help { goptions.PrintHelp() os.Exit(0) } if opts.Version { if Version == "" { fmt.Printf("shield-agent (development)%s\n", Version) } else { fmt.Printf("shield-agent v%s\n", Version) } os.Exit(0) } if opts.ConfigFile == "" { fmt.Fprintf(os.Stderr, "You must specify a configuration file via `--config`\n") os.Exit(1) } log.SetupLogging(log.LogConfig{Type: "console", Level: opts.Log}) log.Infof("starting agent") ag := agent.NewAgent() if err := ag.ReadConfig(opts.ConfigFile); err != nil { log.Errorf("configuration failed: %s", err) return } ag.Run() }
func (oa OAuthenticator) RequireAuth(w http.ResponseWriter, r *http.Request) { sess, err := gothic.Store.Get(r, gothic.SessionName) if err != nil { if _, ok := err.(securecookie.Error); ok { r.Header.Set("Cookie", "") sess, err = gothic.Store.New(r, gothic.SessionName) if err != nil { w.WriteHeader(500) w.Write([]byte("Failure generating a new session: " + err.Error())) return } } else { log.Errorf("%s", err) w.WriteHeader(500) w.Write([]byte("Unexpected error retrieving session data")) return } } sess.AddFlash(r.URL.Path) sess.Save(r, w) // only start oauth redirection if we're hitting the auth APIs, or web UI if ShouldOAuthRedirect(r.URL.Path) { log.Debugf("Starting OAuth Process for request: %s", r) gothic.BeginAuthHandler(w, r) } else { // otherwise set auth header for api clients to understand oauth is needed log.Debugf("Unauthenticated API Request received, OAuth required, sending 401") w.Header().Set("WWW-Authenticate", "Bearer") w.WriteHeader(401) w.Write([]byte("Unauthorized")) } }
func main() { supervisor.Version = Version var opts ShielddOpts opts.Log = "Info" if err := goptions.Parse(&opts); err != nil { fmt.Printf("%s\n", err) goptions.PrintHelp() return } if opts.Help { goptions.PrintHelp() os.Exit(0) } if opts.Version { if Version == "" { fmt.Printf("shieldd (development)\n") } else { fmt.Printf("shieldd v%s\n", Version) } os.Exit(0) } if opts.ConfigFile == "" { fmt.Fprintf(os.Stderr, "No config specified. Please try again using the -c/--config argument\n") os.Exit(1) } log.SetupLogging(log.LogConfig{Type: "console", Level: opts.Log}) log.Infof("starting shield daemon") s := supervisor.NewSupervisor() if err := s.ReadConfig(opts.ConfigFile); err != nil { log.Errorf("Failed to load config: %s", err) return } s.SpawnAPI() s.SpawnWorkers() if err := s.Run(); err != nil { log.Errorf("shield daemon failed: %s", err) } log.Infof("stopping daemon") }
func (agent *Agent) ServeOne(l net.Listener, async bool) { c, err := l.Accept() if err != nil { log.Errorf("failed to accept: %s\n", err) return } conn, chans, reqs, err := ssh.NewServerConn(c, agent.config) if err != nil { log.Errorf("handshake failed: %s\n", err) return } if async { go agent.handleConn(conn, chans, reqs) } else { agent.handleConn(conn, chans, reqs) } }
func (ws *WebServer) Start() { err := ws.Setup() if err != nil { panic("Could not set up WebServer for SHIELD: " + err.Error()) } log.Debugf("Starting WebServer on '%s'...", ws.Addr) err = http.ListenAndServe(ws.Addr, nil) if err != nil { log.Errorf("HTTP API failed %s", err.Error()) panic("Cannot setup WebServer, aborting. Check logs for details.") } }
func JSON(w http.ResponseWriter, thing interface{}) { bytes, err := json.Marshal(thing) if err != nil { log.Errorf("Cannot marshal JSON: <%s>\n", err) w.WriteHeader(500) return } w.Header().Set("Content-Type", "application/json") w.WriteHeader(200) w.Write(bytes) return }
func (s *Supervisor) CheckSchedule() { for _, job := range s.jobq { if !job.Runnable() { continue } log.Infof("scheduling execution of job %s [%s]", job.Name, job.UUID) task, err := s.Database.CreateBackupTask("system", job) if err != nil { log.Errorf("job -> task conversion / database update failed: %s", err) continue } s.ScheduleTask(task) err = job.Reschedule() if err != nil { log.Errorf("error encountered while determining next run of %s (%s): %s", job.UUID, job.Spec, err) } else { log.Infof("next run of %s [%s] which runs %s is at %s", job.Name, job.UUID, job.Spec, job.NextRun) } } }
func main() { var opts ShieldAgentOpts opts.Log = "Info" if err := goptions.Parse(&opts); err != nil { fmt.Printf("%s\n", err) goptions.PrintHelp() return } log.SetupLogging(log.LogConfig{Type: "console", Level: opts.Log}) log.Infof("starting agent") ag := agent.NewAgent() if err := ag.ReadConfig(opts.ConfigFile); err != nil { log.Errorf("configuration failed: %s", err) return } ag.Run() }
func (s *Supervisor) Resync() error { jobq, err := s.Database.GetAllJobs(nil) if err != nil { return err } // calculate the initial run of each job for _, job := range jobq { err := job.Reschedule() if err != nil { log.Errorf("error encountered while determining next run of %s [%s] which runs %s: %s", job.Name, job.UUID, job.Spec, err) } else { log.Infof("initial run of %s [%s] which runs %s is at %s", job.Name, job.UUID, job.Spec, job.NextRun) } } s.jobq = jobq return nil }
func (s *Supervisor) FailUnfinishedTasks() error { tasks, err := s.Database.GetAllTasks( &db.TaskFilter{ ForStatus: db.RunningStatus, }, ) if err != nil { return fmt.Errorf("Failed to sweep database of running tasks: %s", err) } now := time.Now() for _, task := range tasks { log.Warnf("Found task %s in 'running' state at startup; setting to 'failed'", task.UUID) if err := s.Database.FailTask(task.UUID, now); err != nil { return fmt.Errorf("Failed to sweep database of running tasks [%s]: %s", task.UUID, err) } if task.Op == db.BackupOperation && task.ArchiveUUID != nil { archive, err := s.Database.GetArchive(task.ArchiveUUID) if err != nil { log.Warnf("Unable to retrieve archive %s (for task %s) from the database: %s", task.ArchiveUUID, task.UUID, err) continue } log.Warnf("Found archive %s for task %s, purging", archive.UUID, task.UUID) task, err := s.Database.CreatePurgeTask("", archive, s.PurgeAgent) if err != nil { log.Errorf("Failed to purge archive %s (for task %s, which was running at boot): %s", archive.UUID, task.UUID, err) } else { s.ScheduleTask(task) } } } return nil }
func (s *Supervisor) Run() error { if err := s.Database.Connect(); err != nil { return fmt.Errorf("failed to connect to %s database at %s: %s\n", s.Database.Driver, s.Database.DSN, err) } if err := s.Database.CheckCurrentSchema(); err != nil { return fmt.Errorf("database failed schema version check: %s\n", err) } if err := s.Resync(); err != nil { return err } if err := s.FailUnfinishedTasks(); err != nil { return err } if err := s.ReschedulePendingTasks(); err != nil { return err } for { select { case <-s.resync: if err := s.Resync(); err != nil { log.Errorf("resync error: %s", err) } case <-s.purge.C: s.PurgeArchives() case <-s.tick.C: s.CheckSchedule() // see if any tasks have been running past the timeout period if len(s.runq) > 0 { ok := true lst := make([]*db.Task, 0) now := timestamp.Now() for _, runtask := range s.runq { if now.After(runtask.TimeoutAt) { s.Database.CancelTask(runtask.UUID, now.Time()) log.Errorf("shield timed out task '%s' after running for %v", runtask.UUID, s.Timeout) ok = false } else { lst = append(lst, runtask) } } if !ok { s.runq = lst } } // see if we have anything in the schedule queue SchedQueue: for len(s.schedq) > 0 { select { case s.workers <- s.schedq[0]: s.Database.StartTask(s.schedq[0].UUID, time.Now()) s.schedq[0].Attempts++ log.Infof("sent a task to a worker") s.runq = append(s.runq, s.schedq[0]) log.Debugf("added task to the runq") s.schedq = s.schedq[1:] default: break SchedQueue } } case adhoc := <-s.adhoc: s.ScheduleAdhoc(adhoc) case u := <-s.updates: switch u.Op { case STOPPED: log.Infof(" %s: job stopped at %s", u.Task, u.StoppedAt) s.RemoveTaskFromRunq(u.Task) if err := s.Database.CompleteTask(u.Task, u.StoppedAt); err != nil { log.Errorf(" %s: !! failed to update database - %s", u.Task, err) } case FAILED: log.Warnf(" %s: task failed!", u.Task) s.RemoveTaskFromRunq(u.Task) if err := s.Database.FailTask(u.Task, u.StoppedAt); err != nil { log.Errorf(" %s: !! failed to update database - %s", u.Task, err) } case OUTPUT: log.Infof(" %s> %s", u.Task, strings.Trim(u.Output, "\n")) if err := s.Database.UpdateTaskLog(u.Task, u.Output); err != nil { log.Errorf(" %s: !! failed to update database - %s", u.Task, err) } case RESTORE_KEY: log.Infof(" %s: restore key is %s", u.Task, u.Output) if id, err := s.Database.CreateTaskArchive(u.Task, u.Output, time.Now()); err != nil { log.Errorf(" %s: !! failed to update database - %s", u.Task, err) } else { if !u.TaskSuccess { s.Database.InvalidateArchive(id) } } case PURGE_ARCHIVE: log.Infof(" %s: archive %s purged from storage", u.Task, u.Archive) if err := s.Database.PurgeArchive(u.Archive); err != nil { log.Errorf(" %s: !! failed to update database - %s", u.Task, err) } default: log.Errorf(" %s: !! unrecognized op type", u.Task) } } } }
func (s *Supervisor) ScheduleAdhoc(a *db.Task) { log.Infof("schedule adhoc %s job", a.Op) switch a.Op { case db.BackupOperation: // expect a JobUUID to move to the schedq Immediately for _, job := range s.jobq { if !uuid.Equal(job.UUID, a.JobUUID) { continue } log.Infof("scheduling immediate (ad hoc) execution of job %s [%s]", job.Name, job.UUID) task, err := s.Database.CreateBackupTask(a.Owner, job) if err != nil { log.Errorf("job -> task conversion / database update failed: %s", err) if a.TaskUUIDChan != nil { a.TaskUUIDChan <- &db.TaskInfo{ Err: true, Info: err.Error(), } } continue } if a.TaskUUIDChan != nil { a.TaskUUIDChan <- &db.TaskInfo{ Err: false, Info: task.UUID.String(), } } s.ScheduleTask(task) } case db.RestoreOperation: archive, err := s.Database.GetArchive(a.ArchiveUUID) if err != nil { log.Errorf("unable to find archive %s for restore task: %s", a.ArchiveUUID, err) return } target, err := s.Database.GetTarget(a.TargetUUID) if err != nil { log.Errorf("unable to find target %s for restore task: %s", a.TargetUUID, err) return } task, err := s.Database.CreateRestoreTask(a.Owner, archive, target) if err != nil { log.Errorf("restore task database creation failed: %s", err) if a.TaskUUIDChan != nil { a.TaskUUIDChan <- &db.TaskInfo{ Err: true, Info: err.Error(), } } return } if a.TaskUUIDChan != nil { a.TaskUUIDChan <- &db.TaskInfo{ Err: false, Info: task.UUID.String(), } } s.ScheduleTask(task) } }
func (ws *WebServer) Setup() error { var err error log.Debugf("Configuring WebServer...") if err := ws.Database.Connect(); err != nil { log.Errorf("Failed to connect to %s database at %s: %s", ws.Database.Driver, ws.Database.DSN, err) return err } if ws.Auth.OAuth.Provider != "" { log.Debugf("Configuring OAuth Session store") maxSessionAge := ws.Auth.OAuth.Sessions.MaxAge authKey := securecookie.GenerateRandomKey(64) encKey := securecookie.GenerateRandomKey(32) switch ws.Auth.OAuth.Sessions.Type { case "sqlite3": log.Debugf("Using sqlite3 as a session store") store, err := sqlitestore.NewSqliteStore(ws.Auth.OAuth.Sessions.DSN, "http_sessions", "/", maxSessionAge, authKey, encKey) if err != nil { log.Errorf("Error setting up sessions database: %s", err) return err } gothic.Store = store case "postgres": log.Debugf("Using postgres as a session store") gothic.Store = pgstore.NewPGStore(ws.Auth.OAuth.Sessions.DSN, authKey, encKey) gothic.Store.(*pgstore.PGStore).Options.MaxAge = maxSessionAge case "mock": log.Debugf("Using mocked session store") // does nothing, to avoid being accidentally used in prod default: log.Errorf("Invalid DB Backend for OAuth sessions database") return err } gob.Register(map[string]interface{}{}) switch ws.Auth.OAuth.Provider { case "github": log.Debugf("Using github as the oauth provider") goth.UseProviders(github.New(ws.Auth.OAuth.Key, ws.Auth.OAuth.Secret, fmt.Sprintf("%s/v1/auth/github/callback", ws.Auth.OAuth.BaseURL), "read:org")) OAuthVerifier = &GithubVerifier{Orgs: ws.Auth.OAuth.Authorization.Orgs} case "cloudfoundry": log.Debugf("Using cloudfoundry as the oauth provider") goth.UseProviders(cloudfoundry.New(ws.Auth.OAuth.ProviderURL, ws.Auth.OAuth.Key, ws.Auth.OAuth.Secret, fmt.Sprintf("%s/v1/auth/cloudfoundry/callback", ws.Auth.OAuth.BaseURL), "openid,scim.read")) OAuthVerifier = &UAAVerifier{Groups: ws.Auth.OAuth.Authorization.Orgs, UAA: ws.Auth.OAuth.ProviderURL} p, err := goth.GetProvider("cloudfoundry") if err != nil { return err } p.(*cloudfoundry.Provider).Client = ws.Auth.OAuth.Client case "faux": log.Debugf("Using mocked session store") // does nothing, to avoid being accidentally used in prod default: log.Errorf("Invalid OAuth provider specified.") return err } gothic.GetProviderName = func(req *http.Request) (string, error) { return ws.Auth.OAuth.Provider, nil } gothic.SetState = func(req *http.Request) string { sess, _ := gothic.Store.Get(req, gothic.SessionName) sess.Values["state"] = uuid.New() return sess.Values["state"].(string) } } protectedAPIs, err := ws.ProtectedAPIs() if err != nil { log.Errorf("Could not set up HTTP routes: " + err.Error()) return err } if ws.Auth.OAuth.Provider != "" { log.Debugf("Enabling OAuth handlers for HTTP requests") UserAuthenticator = OAuthenticator{ Cfg: ws.Auth.OAuth, } } else { log.Debugf("Enabling Basic Auth handlers for HTTP requests") UserAuthenticator = BasicAuthenticator{ Cfg: ws.Auth.Basic, } } http.Handle("/", ws.UnauthenticatedResources(Authenticate(ws.Auth.Tokens, protectedAPIs))) return nil }
func (oa OAuthenticator) OAuthCallback() http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { log.Debugf("Incoming Auth request: %s", r) sess, err := gothic.Store.Get(r, gothic.SessionName) if err != nil { log.Errorf("Error retrieving session info: %s", err) w.WriteHeader(500) return } log.Debugf("Processing oauth callback for '%s'", sess.ID) if gothic.GetState(r) != sess.Values["state"] { w.WriteHeader(403) w.Write([]byte("Unauthorized")) return } if r.URL.Query().Get("code") == "" { log.Errorf("No code detected in oauth callback: %v", r) w.WriteHeader(403) w.Write([]byte("No oauth code issued from provider")) return } user, err := gothic.CompleteUserAuth(w, r) if err != nil { log.Errorf("Error verifying oauth success: %s. Request: %v", err, r) w.WriteHeader(403) w.Write([]byte("UnOAuthorized")) return } log.Debugf("Authenticated user %#v", user) ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: user.AccessToken}) ctx := context.WithValue(oauth2.NoContext, oauth2.HTTPClient, oa.Cfg.Client) tc := oauth2.NewClient(ctx, ts) log.Debugf("Checking authorization...") membership, err := OAuthVerifier.Membership(user, tc) if err != nil { log.Errorf("Error retreiving user membership: %s", err) w.WriteHeader(403) w.Write([]byte("Unable to verify your membership")) return } if !OAuthVerifier.Verify(user.NickName, membership) { log.Debugf("Authorization denied") w.WriteHeader(403) w.Write([]byte("You are not authorized to view this content")) return } log.Infof("Successful login for %s", user.NickName) redirect := "/" if flashes := sess.Flashes(); len(flashes) > 0 { if flash, ok := flashes[0].(string); ok { // don't redirect back to api calls, to prevent auth redirection loops if !apiCall.MatchString(flash) || cliAuthCall.MatchString(flash) { redirect = flash } } } sess.Values["User"] = user.NickName sess.Values["Membership"] = membership err = sess.Save(r, w) if err != nil { log.Errorf("Error saving session: %s", err) w.WriteHeader(500) w.Write([]byte("Unable to save authentication data. Check the SHIELD logs for more info.")) return } http.Redirect(w, r, redirect, 302) // checks auth }) }
func bail(w http.ResponseWriter, e error) { w.WriteHeader(500) log.Errorf("Request bailed: <%s>\n", e) return }
func worker(id uint, privateKeyFile string, work chan *db.Task, updates chan WorkerUpdate) { config, err := agent.ConfigureSSHClient(privateKeyFile) if err != nil { log.Errorf("worker %d unable to read user key %s: %s; bailing out.\n", id, privateKeyFile, err) return } for t := range work { client := agent.NewClient(config) remote := t.Agent if remote == "" { updates <- WorkerUpdate{Task: t.UUID, Op: OUTPUT, Output: fmt.Sprintf("TASK FAILED!! no remote agent specified for task %s\n", t.UUID)} updates <- WorkerUpdate{Task: t.UUID, Op: FAILED, StoppedAt: time.Now()} continue } err = client.Dial(remote) if err != nil { updates <- WorkerUpdate{Task: t.UUID, Op: OUTPUT, Output: fmt.Sprintf("TASK FAILED!! shield worker %d unable to connect to %s (%s)\n", id, remote, err)} updates <- WorkerUpdate{Task: t.UUID, Op: FAILED, StoppedAt: time.Now()} continue } // start a command and stream output final := make(chan string) partial := make(chan string) go func(out chan string, up chan WorkerUpdate, id uuid.UUID, in chan string) { var buffer []string for { s, ok := <-in if !ok { break } switch s[0:2] { case "O:": buffer = append(buffer, s[2:]) case "E:": up <- WorkerUpdate{ Task: id, Op: OUTPUT, Output: s[2:] + "\n", } } } out <- strings.Join(buffer, "") close(out) }(final, updates, t.UUID, partial) command, err := json.Marshal(WorkerRequest{ Operation: t.Op, TargetPlugin: t.TargetPlugin, TargetEndpoint: t.TargetEndpoint, StorePlugin: t.StorePlugin, StoreEndpoint: t.StoreEndpoint, RestoreKey: t.RestoreKey, }) if err != nil { updates <- WorkerUpdate{Task: t.UUID, Op: OUTPUT, Output: fmt.Sprintf("TASK FAILED!! shield worker %d was unable to json encode the request bound for remote agent %s (%s)", id, remote, err), } updates <- WorkerUpdate{Task: t.UUID, Op: FAILED, StoppedAt: time.Now()} client.Close() continue } // exec the command var jobFailed bool err = client.Run(partial, string(command)) if err != nil { updates <- WorkerUpdate{Task: t.UUID, Op: OUTPUT, Output: fmt.Sprintf("TASK FAILED!! shield worker %d failed to execute the command against the remote agent %s (%s)\n", id, remote, err)} jobFailed = true } client.Close() out := <-final if t.Op == db.BackupOperation { // parse JSON from standard output and get the restore key // (this might fail, we might not get a key, etc.) v := struct { Key string }{} buf := bytes.NewBufferString(out) dec := json.NewDecoder(buf) err := dec.Decode(&v) if err != nil { jobFailed = true updates <- WorkerUpdate{Task: t.UUID, Op: OUTPUT, Output: fmt.Sprintf("WORKER FAILED!! shield worker %d failed to parse JSON response from remote agent %s (%s)\n", id, remote, err)} } else { if v.Key != "" { updates <- WorkerUpdate{ Task: t.UUID, Op: RESTORE_KEY, TaskSuccess: !jobFailed, Output: v.Key, } } else { jobFailed = true updates <- WorkerUpdate{Task: t.UUID, Op: OUTPUT, Output: fmt.Sprintf("TASK FAILED!! No restore key detected in worker %d. Cowardly refusing to create an archive record", id)} } } } if t.Op == db.PurgeOperation && !jobFailed { updates <- WorkerUpdate{ Task: t.UUID, Op: PURGE_ARCHIVE, Archive: t.ArchiveUUID, } } // signal to the supervisor that we finished if jobFailed { updates <- WorkerUpdate{Task: t.UUID, Op: FAILED, StoppedAt: time.Now()} } else { updates <- WorkerUpdate{ Task: t.UUID, Op: STOPPED, StoppedAt: time.Now(), } } } }
func (agent *Agent) handleConn(conn *ssh.ServerConn, chans <-chan ssh.NewChannel, reqs <-chan *ssh.Request) { defer conn.Close() for newChannel := range chans { if newChannel.ChannelType() != "session" { log.Errorf("rejecting unknown channel type: %s\n", newChannel.ChannelType()) newChannel.Reject(ssh.UnknownChannelType, "unknown channel type") continue } channel, requests, err := newChannel.Accept() if err != nil { log.Errorf("failed to accept channel: %s\n", err) return } defer channel.Close() for req := range requests { if req.Type != "exec" { log.Errorf("rejecting non-exec channel request (type=%s)\n", req.Type) req.Reply(false, nil) continue } request, err := ParseRequest(req) if err != nil { log.Errorf("%s\n", err) req.Reply(false, nil) continue } if err = request.ResolvePaths(agent); err != nil { log.Errorf("%s\n", err) req.Reply(false, nil) continue } //log.Errorf("got an agent-request [%s]\n", request.JSON) req.Reply(true, nil) // drain output to the SSH channel stream output := make(chan string) done := make(chan int) go func(out io.Writer, in chan string, done chan int) { for { s, ok := <-in if !ok { break } fmt.Fprintf(out, "%s", s) log.Debugf("%s", s) } close(done) }(channel, output, done) // run the agent request err = request.Run(output) <-done var rc int if exitErr, ok := err.(*exec.ExitError); ok { sys := exitErr.ProcessState.Sys() // os.ProcessState.Sys() may not return syscall.WaitStatus on non-UNIX machines, // so currently this feature only works on UNIX, but shouldn't crash on other OSes if ws, ok := sys.(syscall.WaitStatus); ok { if ws.Exited() { rc = ws.ExitStatus() } else { var signal syscall.Signal if ws.Signaled() { signal = ws.Signal() } if ws.Stopped() { signal = ws.StopSignal() } sigStr, ok := SIGSTRING[signal] if !ok { sigStr = "ABRT" // use ABRT as catch-all signal for any that don't translate log.Infof("Task execution terminted due to %s, translating as ABRT for ssh transport", signal) } else { log.Infof("Task execution terminated due to SIG%s", sigStr) } sigMsg := struct { Signal string CoreDumped bool Error string Lang string }{ Signal: sigStr, CoreDumped: false, Error: fmt.Sprintf("shield-pipe terminated due to SIG%s", sigStr), Lang: "en-US", } channel.SendRequest("exit-signal", false, ssh.Marshal(&sigMsg)) channel.Close() continue } } } else if err != nil { // we got some kind of error that isn't a command execution error, // from a UNIX system, use an magical error code to signal this to // the shield daemon - 16777216 log.Infof("Task could not execute: %s", err) rc = 16777216 } log.Infof("Task completed with rc=%d", rc) byteCode := make([]byte, 4) binary.BigEndian.PutUint32(byteCode, uint32(rc)) // SSH protocol is big-endian byte ordering channel.SendRequest("exit-status", false, byteCode) channel.Close() } } }