github.com/mattyw/juju@v0.0.0-20140610034352-732aecd63861/worker/authenticationworker/worker.go (about)

     1  // Copyright 2013 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package authenticationworker
     5  
     6  import (
     7  	"strings"
     8  
     9  	"github.com/juju/errors"
    10  	"github.com/juju/loggo"
    11  	"github.com/juju/utils/set"
    12  	"launchpad.net/tomb"
    13  
    14  	"github.com/juju/juju/agent"
    15  	"github.com/juju/juju/state/api/keyupdater"
    16  	"github.com/juju/juju/state/api/watcher"
    17  	"github.com/juju/juju/utils/ssh"
    18  	"github.com/juju/juju/worker"
    19  )
    20  
    21  // The user name used to ssh into Juju nodes.
    22  // Override for testing.
    23  var SSHUser = "ubuntu"
    24  
    25  var logger = loggo.GetLogger("juju.worker.authenticationworker")
    26  
    27  type keyupdaterWorker struct {
    28  	st   *keyupdater.State
    29  	tomb tomb.Tomb
    30  	tag  string
    31  	// jujuKeys are the most recently retrieved keys from state.
    32  	jujuKeys set.Strings
    33  	// nonJujuKeys are those added externally to auth keys file
    34  	// such keys do not have comments with the Juju: prefix.
    35  	nonJujuKeys []string
    36  }
    37  
    38  var _ worker.NotifyWatchHandler = (*keyupdaterWorker)(nil)
    39  
    40  // NewWorker returns a worker that keeps track of
    41  // the machine's authorised ssh keys and ensures the
    42  // ~/.ssh/authorized_keys file is up to date.
    43  func NewWorker(st *keyupdater.State, agentConfig agent.Config) worker.Worker {
    44  	kw := &keyupdaterWorker{st: st, tag: agentConfig.Tag()}
    45  	return worker.NewNotifyWorker(kw)
    46  }
    47  
    48  // SetUp is defined on the worker.NotifyWatchHandler interface.
    49  func (kw *keyupdaterWorker) SetUp() (watcher.NotifyWatcher, error) {
    50  	// Record the keys Juju knows about.
    51  	jujuKeys, err := kw.st.AuthorisedKeys(kw.tag)
    52  	if err != nil {
    53  		return nil, errors.LoggedErrorf(logger, "reading Juju ssh keys for %q: %v", kw.tag, err)
    54  	}
    55  	kw.jujuKeys = set.NewStrings(jujuKeys...)
    56  
    57  	// Read the keys currently in ~/.ssh/authorised_keys.
    58  	sshKeys, err := ssh.ListKeys(SSHUser, ssh.FullKeys)
    59  	if err != nil {
    60  		return nil, errors.LoggedErrorf(logger, "reading ssh authorized keys for %q: %v", kw.tag, err)
    61  	}
    62  	// Record any keys not added by Juju.
    63  	for _, key := range sshKeys {
    64  		_, comment, err := ssh.KeyFingerprint(key)
    65  		// Also record keys which we cannot parse.
    66  		if err != nil || !strings.HasPrefix(comment, ssh.JujuCommentPrefix) {
    67  			kw.nonJujuKeys = append(kw.nonJujuKeys, key)
    68  		}
    69  	}
    70  	// Write out the ssh authorised keys file to match the current state of the world.
    71  	if err := kw.writeSSHKeys(jujuKeys); err != nil {
    72  		return nil, errors.LoggedErrorf(logger, "adding current Juju keys to ssh authorised keys: %v", err)
    73  	}
    74  
    75  	w, err := kw.st.WatchAuthorisedKeys(kw.tag)
    76  	if err != nil {
    77  		return nil, errors.LoggedErrorf(logger, "starting key updater worker: %v", err)
    78  	}
    79  	logger.Infof("%q key updater worker started", kw.tag)
    80  	return w, nil
    81  }
    82  
    83  // writeSSHKeys writes out a new ~/.ssh/authorised_keys file, retaining any non Juju keys
    84  // and adding the specified set of Juju keys.
    85  func (kw *keyupdaterWorker) writeSSHKeys(jujuKeys []string) error {
    86  	allKeys := kw.nonJujuKeys
    87  	// Ensure any Juju keys have the required prefix in their comment.
    88  	for i, key := range jujuKeys {
    89  		jujuKeys[i] = ssh.EnsureJujuComment(key)
    90  	}
    91  	allKeys = append(allKeys, jujuKeys...)
    92  	return ssh.ReplaceKeys(SSHUser, allKeys...)
    93  }
    94  
    95  // Handle is defined on the worker.NotifyWatchHandler interface.
    96  func (kw *keyupdaterWorker) Handle() error {
    97  	// Read the keys that Juju has.
    98  	newKeys, err := kw.st.AuthorisedKeys(kw.tag)
    99  	if err != nil {
   100  		return errors.LoggedErrorf(logger, "reading Juju ssh keys for %q: %v", kw.tag, err)
   101  	}
   102  	// Figure out if any keys have been added or deleted.
   103  	newJujuKeys := set.NewStrings(newKeys...)
   104  	deleted := kw.jujuKeys.Difference(newJujuKeys)
   105  	added := newJujuKeys.Difference(kw.jujuKeys)
   106  	if added.Size() > 0 || deleted.Size() > 0 {
   107  		logger.Debugf("adding ssh keys to authorised keys: %v", added)
   108  		logger.Debugf("deleting ssh keys from authorised keys: %v", deleted)
   109  		if err = kw.writeSSHKeys(newKeys); err != nil {
   110  			return errors.LoggedErrorf(logger, "updating ssh keys: %v", err)
   111  		}
   112  	}
   113  	kw.jujuKeys = newJujuKeys
   114  	return nil
   115  }
   116  
   117  // TearDown is defined on the worker.NotifyWatchHandler interface.
   118  func (kw *keyupdaterWorker) TearDown() error {
   119  	// Nothing to do here.
   120  	return nil
   121  }