import crochet import pappyproxy import shlex from pappyproxy.plugin import main_context from pappyproxy.console import load_reqlist from pappyproxy.util import PappyException from twisted.internet import defer @crochet.wait_for(timeout=None) @defer.inlineCallbacks def tag(line): """ Add a tag to requests. Usage: tag [request ids] You can tag as many requests as you want at the same time. If no ids are given, the tag will be applied to all in-context requests. """ args = shlex.split(line) if len(args) == 0: raise PappyException('Tag name is required') tag = args[0] if len(args) > 1: reqs = yield load_reqlist(args[1], False) ids = [r.reqid for r in reqs] print 'Tagging %s with %s' % (', '.join(ids), tag) else: print "Tagging all in-context requests with %s" % tag reqs = main_context().active_requests for req in reqs: if tag not in req.tags: req.tags.append(tag) if req.saved: yield req.async_save() add_req(req) else: print 'Request %s already has tag %s' % (req.reqid, tag) @crochet.wait_for(timeout=None) @defer.inlineCallbacks def untag(line): """ Remove a tag from requests Usage: untag You can provide as many request ids as you want and the tag will be removed from all of them. If no ids are given, the tag will be removed from all in-context requests. """ args = shlex.split(line) if len(args) == 0: raise PappyException("Tag and request ids are required") tag = args[0] ids = [] if len(args) > 1: reqs = yield load_reqlist(args[1], False) ids = [r.reqid for r in reqs] else: print "Untagging all in-context requests with tag %s" % tag reqs = main_context().active_requests for req in reqs: if tag in req.tags: req.tags.remove(tag) if req.saved: yield req.async_save() if ids: print 'Tag %s removed from %s' % (tag, ', '.join(ids)) @crochet.wait_for(timeout=None) @defer.inlineCallbacks def clrtag(line): """ Clear all the tags from requests Usage: clrtag """ args = shlex.split(line) if len(args) == 0: raise PappyException('No request IDs given') reqs = yield load_reqlist(args[0], False) for req in reqs: if req.tags: req.tags = [] print 'Tags cleared from request %s' % (req.reqid) if req.saved: yield req.async_save() ############### ## Plugin hooks def load_cmds(cmd): cmd.set_cmds({ 'clrtag': (clrtag, None), 'untag': (untag, None), 'tag': (tag, None), }) cmd.add_aliases([ #('rpy', ''), ])