/client/actions.py
Python | 182 lines | 97 code | 25 blank | 60 comment | 10 complexity | 302b41dfb7aebbb12c97cc449fe36a78 MD5 | raw file
Possible License(s): Apache-2.0
- #!/usr/bin/env python
- # Copyright 2010 Google 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.
- """This file contains common grr jobs."""
- import logging
- import pdb
- import traceback
- from grr.client import conf as flags
- from grr.lib import registry
- from grr.lib import utils
- from grr.proto import jobs_pb2
- FLAGS = flags.FLAGS
- # Our first response in the session is this:
- INITIAL_RESPONSE_ID = 1
- class ActionPlugin(object):
- """Baseclass for plugins.
- An action is a plugin abstraction which receives a protobuf and
- sends a protobuf in response.
- The code is specified in the Run() method, while the data is
- specified in the in_protobuf and out_protobuf classes.
- """
- # The protobuf that will be used to encode this message
- in_protobuf = None
- # The protobuf type we send
- out_protobuf = None
- # Authentication Required for this Action:
- _authentication_required = True
- __metaclass__ = registry.MetaclassRegistry
- def __init__(self, message, grr_context=None, **proto_args):
- """Initialises our protobuf from the keywords passed.
- Args:
- message: The GrrMessage that we are called to process.
- grr_context: The grr context object which may be used to
- e.g. send new actions on.
- **proto_args: Field initializers for the protobuf in self._protobuf.
- """
- self.grr_context = grr_context
- self.message = message
- self.response_id = INITIAL_RESPONSE_ID
- if self.in_protobuf:
- self.buff = self.in_protobuf()
- for k, v in proto_args.items():
- setattr(self.buff, k, v)
- def Execute(self, message):
- """This function parses the protobuf from the server.
- The Run method will be called with the unserialised protobuf.
- Args:
- message: The encoded protobuf which will be decoded
- by the plugin.
- Returns:
- Upon return a callback will be called on the server to register
- the end of the function and pass back exceptions.
- """
- if self.in_protobuf:
- args = self.in_protobuf()
- args.ParseFromString(message.args)
- else:
- args = None
- self.status = jobs_pb2.GrrStatus()
- self.status.status = jobs_pb2.GrrStatus.OK # Default status.
- try:
- # Only allow authenticated messages in the client
- if (self._authentication_required and
- message.auth_state != jobs_pb2.GrrMessage.AUTHENTICATED):
- raise RuntimeError("Message for %s was not Authenticated." %
- message.name)
- self.Run(args)
- # We want to report back all errors and map Python exceptions to
- # Grr Errors.
- except Exception, e:
- self.SetStatus(jobs_pb2.GrrStatus.GENERIC_ERROR, "%r: %s" % (e, e),
- traceback.format_exc())
- if FLAGS.debug:
- pdb.post_mortem()
- if self.status.status != jobs_pb2.GrrStatus.OK:
- logging.info("Job Error (%s): %s", self.__class__.__name__,
- self.status.error_message)
- if self.status.backtrace:
- logging.debug(self.status.backtrace)
- # This returns the error status of the Actions to the flow.
- self.SendReply(self.status, message_type=jobs_pb2.GrrMessage.STATUS)
- def Run(self, unused_args):
- """Main plugin entry point.
- This function will always be overridden by real plugins.
- Args:
- unused_args: An already initialised protobuf object.
- Raises:
- KeyError: if not implemented.
- """
- raise KeyError("Action %s not available on this platform." %
- self.message.name)
- def SetStatus(self, status, message="", backtrace=None):
- """Set a status to report back to the server."""
- self.status.status = status
- self.status.error_message = utils.SmartUnicode(message)
- if backtrace:
- self.status.backtrace = utils.SmartUnicode(backtrace)
- def SendReply(self, protobuf=None, message_type=jobs_pb2.GrrMessage.MESSAGE,
- **kw):
- """Send response back to the server."""
- if protobuf is None:
- protobuf = self.out_protobuf(**kw)
- self.grr_context.SendReply(protobuf,
- session_id=self.message.session_id,
- response_id=self.response_id,
- request_id=self.message.request_id,
- message_type=message_type)
- self.response_id += 1
- class IteratedAction(ActionPlugin):
- """An action which can restore its state from an iterator.
- Implement iterating actions by extending this class and overriding the
- Iterate() method.
- """
- def Run(self, request):
- """Munge the iterator to the server and abstract it away."""
- # Pass the client_state as a dict to the action. This is often more
- # efficient than manipulating a protobuf.
- client_state = utils.ProtoDict(request.iterator.client_state).ToDict()
- # Derived classes should implement this.
- self.Iterate(request, client_state)
- # Update the iterator client_state from the dict.
- request.iterator.client_state.CopyFrom(
- utils.ProtoDict(client_state).ToProto())
- # Return the iterator
- self.SendReply(request.iterator, message_type=jobs_pb2.GrrMessage.ITERATOR)
- def Iterate(self, request, client_state):
- """Actions should override this."""