Update testing scripts.

3 files changed, 325 insertions, 121 deletions

diff --git a/scripts/db_access/reset_db.py b/scripts/db_access/reset_db.py
index 80871856..7c6d281b 100644
--- a/scripts/db_access/reset_db.py
+++ b/scripts/db_access/reset_db.py
@@ -5,20 +5,21 @@
 # WARNING: running this script over a database will delete all documents but
 # the one with id u1db_config (which contains db metadata) and design docs
 # needed for couch backend.
+#
+# Run it like this to get some help:
+#
+#     ./reset_db.py --help
 
 
-import sys
-from ConfigParser import ConfigParser
 import threading
 import logging
-from couchdb import Database as CouchDatabase
-
+import argparse
+import re
 
-if len(sys.argv) != 2:
-    print 'Usage: %s <uuid>' % sys.argv[0]
-    exit(1)
 
-uuid = sys.argv[1]
+from ConfigParser import ConfigParser
+from couchdb import Database as CouchDatabase
+from couchdb import Server as CouchServer
 
 
 # create a logger
@@ -27,23 +28,6 @@ LOG_FORMAT = '%(asctime)s %(message)s'
 logging.basicConfig(format=LOG_FORMAT, level=logging.INFO)
 
 
-# get couch url
-cp = ConfigParser()
-cp.read('/etc/leap/soledad-server.conf')
-url = cp.get('soledad-server', 'couch_url')
-
-
-# confirm
-yes = raw_input("Are you sure you want to reset the database for user %s "
-                "(type YES)? " % uuid)
-if yes != 'YES':
-    print 'Bailing out...'
-    exit(2)
-
-
-db = CouchDatabase('%s/user-%s' % (url, uuid))
-
-
 class _DeleterThread(threading.Thread):
 
     def __init__(self, db, doc_id, release_fun):
@@ -59,21 +43,95 @@ class _DeleterThread(threading.Thread):
         self._release_fun()
 
 
-semaphore_pool = threading.BoundedSemaphore(value=20)
-
-
-threads = []
-for doc_id in db:
-    if doc_id != 'u1db_config' and not doc_id.startswith('_design'):
+def get_confirmation(noconfirm, uuid, shared):
+    msg = "Are you sure you want to reset %s (type YES)? "
+    if shared:
+        msg = msg % "the shared database"
+    elif uuid:
+        msg = msg % ("the database for user %s" % uuid)
+    else:
+        msg = msg % "all databases"
+    if noconfirm is False:
+        yes = raw_input(msg)
+        if yes != 'YES':
+            print 'Bailing out...'
+            exit(2)
+
+
+def get_url(empty):
+    url = None
+    if empty is False:
+        # get couch url
+        cp = ConfigParser()
+        cp.read('/etc/leap/soledad-server.conf')
+        url = cp.get('soledad-server', 'couch_url')
+    else:
+        with open('/etc/couchdb/couchdb.netrc') as f:
+            netrc = f.read()
+            admin_password = re.match('^.* password (.*)$', netrc).groups()[0]
+            url = 'http://admin:%s@127.0.0.1:5984' % admin_password
+    return url
+
+
+def reset_all_dbs(url, empty):
+    server = CouchServer('%s' % (url))
+    for dbname in server:
+        if dbname.startswith('user-') or dbname == 'shared':
+            reset_db(url, dbname, empty)
+
+
+def reset_db(url, dbname, empty):
+    db = CouchDatabase('%s/%s' % (url, dbname))
+    semaphore_pool = threading.BoundedSemaphore(value=20)
+
+    # launch threads for deleting docs
+    threads = []
+    for doc_id in db:
+        if empty is False:
+            if doc_id == 'u1db_config' or doc_id.startswith('_design'):
+                continue
         semaphore_pool.acquire()
         logger.info('[main] launching thread for doc: %s' % doc_id)
         t = _DeleterThread(db, doc_id, semaphore_pool.release)
         t.start()
         threads.append(t)
 
-
-logger.info('[main] waiting for threads.')
-map(lambda thread: thread.join(), threads)
-
-
-logger.info('[main] done.')
+    # wait for threads to finish
+    logger.info('[main] waiting for threads.')
+    map(lambda thread: thread.join(), threads)
+    logger.info('[main] done.')
+
+
+def _parse_args():
+    parser = argparse.ArgumentParser()
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument('-u', dest='uuid', default=False,
+        help='Reset database of given user.')
+    group.add_argument('-s', dest='shared', action='store_true', default=False,
+        help='Reset the shared database.')
+    group.add_argument('-a', dest='all', action='store_true', default=False,
+        help='Reset all user databases.')
+    parser.add_argument(
+        '-e', dest='empty', action='store_true', required=False, default=False,
+        help='Empty database (do not preserve minimal set of u1db documents).')
+    parser.add_argument(
+        '-y', dest='noconfirm', action='store_true', required=False,
+        default=False,
+        help='Do not ask for confirmation.')
+    return parser.parse_args(), parser
+
+
+if __name__ == '__main__':
+    args, parser = _parse_args()
+    if not (args.uuid or args.shared or args.all):
+        parser.print_help()
+        exit(1)
+
+    url = get_url(args.empty)
+    get_confirmation(args.noconfirm, args.uuid, args.shared)
+    if args.uuid:
+        reset_db(url, "user-%s" % args.uuid, args.empty)
+    elif args.shared:
+        reset_db(url, "shared", args.empty)
+    elif args.all:
+        reset_all_dbs(url, args.empty)
diff --git a/scripts/ddocs/update_design_docs.py b/scripts/ddocs/update_design_docs.py
index e7b5a29c..2e2fa8f0 100644
--- a/scripts/ddocs/update_design_docs.py
+++ b/scripts/ddocs/update_design_docs.py
@@ -11,84 +11,83 @@ import re
 import threading
 import binascii
 
-
+from urlparse import urlparse
 from getpass import getpass
 from ConfigParser import ConfigParser
-from couchdb.client import Server
-from couchdb.http import Resource, Session
-from datetime import datetime
-from urlparse import urlparse
 
+from couchdb.client import Server
+from couchdb.http import Resource
+from couchdb.http import Session
+from couchdb.http import ResourceNotFound
 
 from leap.soledad.common import ddocs
 
 
-# parse command line for the log file name
-logger_fname = "/tmp/update-design-docs_%s.log" % \
-               str(datetime.now()).replace(' ', '_')
-parser = argparse.ArgumentParser()
-parser.add_argument('--log', action='store', default=logger_fname, type=str,
-                    required=False, help='the name of the log file', nargs=1)
-args = parser.parse_args()
+MAX_THREADS = 20
+DESIGN_DOCS = {
+    '_design/docs': json.loads(binascii.a2b_base64(ddocs.docs)),
+    '_design/syncs': json.loads(binascii.a2b_base64(ddocs.syncs)),
+    '_design/transactions': json.loads(
+        binascii.a2b_base64(ddocs.transactions)),
+}
 
 
-# configure the logger
+# create a logger
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.DEBUG)
-print "Logging to %s." % args.log
-logging.basicConfig(
-    filename=args.log,
-    format="%(asctime)-15s %(message)s")
+LOG_FORMAT = '%(asctime)s %(message)s'
+logging.basicConfig(format=LOG_FORMAT, level=logging.INFO)
 
 
-# configure threads
-max_threads = 20
-semaphore_pool = threading.BoundedSemaphore(value=max_threads)
-threads = []
+def _parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-u', dest='uuid', default=None, type=str,
+                        help='the UUID of the user')
+    parser.add_argument('-t', dest='threads', default=MAX_THREADS, type=int,
+                        help='the number of parallel threads')
+    return parser.parse_args()
 
-# get couch url
-cp = ConfigParser()
-cp.read('/etc/leap/soledad-server.conf')
-url = urlparse(cp.get('soledad-server', 'couch_url'))
 
-# get admin password
-netloc = re.sub('^.*@', '', url.netloc)
-url = url._replace(netloc=netloc)
-password = getpass("Admin password for %s: " % url.geturl())
-url = url._replace(netloc='admin:%s@%s' % (password, netloc))
+def _get_url():
+    # get couch url
+    cp = ConfigParser()
+    cp.read('/etc/leap/soledad-server.conf')
+    url = urlparse(cp.get('soledad-server', 'couch_url'))
+    # get admin password
+    netloc = re.sub('^.*@', '', url.netloc)
+    url = url._replace(netloc=netloc)
+    password = getpass("Admin password for %s: " % url.geturl())
+    return url._replace(netloc='admin:%s@%s' % (password, netloc))
 
-resource = Resource(url.geturl(), Session(retry_delays=[1,2,4,8], timeout=10))
-server = Server(url=resource)
 
-hidden_url = re.sub(
-    'http://(.*):.*@',
-    'http://\\1:xxxxx@',
-    url.geturl())
+def _get_server(url):
+    resource = Resource(
+        url.geturl(), Session(retry_delays=[1, 2, 4, 8], timeout=10))
+    return Server(url=resource)
 
-print """
-==========
-ATTENTION!
-==========
 
-This script will modify Soledad's shared and user databases in:
+def _confirm(url):
+    hidden_url = re.sub(
+        'http://(.*):.*@',
+        'http://\\1:xxxxx@',
+        url.geturl())
 
-  %s
+    print """
+    ==========
+    ATTENTION!
+    ==========
 
-This script does not make a backup of the couch db data, so make sure you
-have a copy or you may loose data.
-""" % hidden_url
-confirm = raw_input("Proceed (type uppercase YES)? ")
+    This script will modify Soledad's shared and user databases in:
 
-if confirm != "YES":
-    exit(1)
+      %s
 
-# convert design doc content
+    This script does not make a backup of the couch db data, so make sure you
+    have a copy or you may loose data.
+    """ % hidden_url
+    confirm = raw_input("Proceed (type uppercase YES)? ")
+
+    if confirm != "YES":
+        exit(1)
 
-design_docs = {
-    '_design/docs': json.loads(binascii.a2b_base64(ddocs.docs)),
-    '_design/syncs': json.loads(binascii.a2b_base64(ddocs.syncs)),
-    '_design/transactions': json.loads(binascii.a2b_base64(ddocs.transactions)),
-}
 
 #
 # Thread
@@ -106,42 +105,66 @@ class DBWorkerThread(threading.Thread):
 
     def run(self):
 
-        logger.info("(%d/%d) Updating db %s." % (self._db_idx, self._db_len,
-                    self._dbname))
+        logger.info(
+            "(%d/%d) Updating db %s."
+            % (self._db_idx, self._db_len, self._dbname))
 
-        for doc_id in design_docs:
-            doc = self._cdb[doc_id]
+        for doc_id in DESIGN_DOCS:
+            try:
+                doc = self._cdb[doc_id]
+            except ResourceNotFound:
+                doc = {'_id': doc_id}
             for key in ['lists', 'views', 'updates']:
-                if key in design_docs[doc_id]:
-                    doc[key] = design_docs[doc_id][key]
+                if key in DESIGN_DOCS[doc_id]:
+                    doc[key] = DESIGN_DOCS[doc_id][key]
             self._cdb.save(doc)
 
         # release the semaphore
         self._release_fun()
 
 
-db_idx = 0
-db_len = len(server)
-for dbname in server:
-
-    db_idx += 1
-
-    if not (dbname.startswith('user-') or dbname == 'shared') \
-            or dbname == 'user-test-db':
-        logger.info("(%d/%d) Skipping db %s." % (db_idx, db_len, dbname))
-        continue
-
-
-    # get access to couch db
-    cdb = Server(url.geturl())[dbname]
-
-    #---------------------------------------------------------------------
-    # Start DB worker thread
-    #---------------------------------------------------------------------
-    semaphore_pool.acquire()
-    thread = DBWorkerThread(server, dbname, db_idx, db_len, semaphore_pool.release)
+def _launch_update_design_docs_thread(
+        server, dbname, db_idx, db_len, semaphore_pool):
+    semaphore_pool.acquire()  # wait for an available working slot
+    thread = DBWorkerThread(
+        server, dbname, db_idx, db_len, semaphore_pool.release)
     thread.daemon = True
     thread.start()
-    threads.append(thread)
-
-map(lambda thread: thread.join(), threads)
+    return thread
+
+
+def _update_design_docs(args, server):
+
+    # find the actual databases to be updated
+    dbs = []
+    if args.uuid:
+        dbs.append('user-%s' % args.uuid)
+    else:
+        for dbname in server:
+            if dbname.startswith('user-') or dbname == 'shared':
+                dbs.append(dbname)
+            else:
+                logger.info("Skipping db %s." % dbname)
+
+    db_idx = 0
+    db_len = len(dbs)
+    semaphore_pool = threading.BoundedSemaphore(value=args.threads)
+    threads = []
+
+    # launch the update
+    for db in dbs:
+        db_idx += 1
+        threads.append(
+            _launch_update_design_docs_thread(
+                server, db, db_idx, db_len, semaphore_pool))
+
+    # wait for all threads to finish
+    map(lambda thread: thread.join(), threads)
+
+
+if __name__ == "__main__":
+    args = _parse_args()
+    url = _get_url()
+    _confirm(url)
+    server = _get_server(url)
+    _update_design_docs(args, server)
diff --git a/scripts/profiling/spam.py b/scripts/profiling/spam.py
new file mode 100755
index 00000000..091a8c48
--- /dev/null
+++ b/scripts/profiling/spam.py
@@ -0,0 +1,123 @@
+#!/usr/bin/python
+
+# Send a lot of messages in parallel.
+
+
+import string
+import smtplib
+import threading
+import logging
+
+from argparse import ArgumentParser
+
+
+SMTP_HOST = 'chipmonk.cdev.bitmask.net'
+NUMBER_OF_THREADS = 20
+
+
+logger = logging.getLogger(__name__)
+LOG_FORMAT = '%(asctime)s %(message)s'
+logging.basicConfig(format=LOG_FORMAT, level=logging.INFO)
+
+
+def _send_email(host, subject, to_addr, from_addr, body_text):
+    """
+    Send an email
+    """
+    body = string.join((
+            "From: %s" % from_addr,
+            "To: %s" % to_addr,
+            "Subject: %s" % subject,
+            "",
+            body_text
+            ), "\r\n")
+    server = smtplib.SMTP(host)
+    server.sendmail(from_addr, [to_addr], body)
+    server.quit()
+
+
+def _parse_args():
+    parser = ArgumentParser()
+    parser.add_argument(
+        'target_address',
+        help='The target email address to spam')
+    parser.add_argument(
+        'number_of_messages', type=int,
+        help='The amount of messages email address to spam')
+    parser.add_argument(
+        '-s', dest='server', default=SMTP_HOST,
+        help='The SMTP server to use')
+    parser.add_argument(
+        '-t', dest='threads', default=NUMBER_OF_THREADS,
+        help='The maximum number of parallel threads to launch')
+    return parser.parse_args()
+
+
+class EmailSenderThread(threading.Thread):
+
+    def __init__(self, host, subject, to_addr, from_addr, body_text,
+            finished_fun):
+        threading.Thread.__init__(self)
+        self._host = host
+        self._subject = subject
+        self._to_addr = to_addr
+        self._from_addr = from_addr
+        self._body_text = body_text
+        self._finished_fun = finished_fun
+
+    def run(self):
+        _send_email(
+            self._host, self._subject, self._to_addr, self._from_addr,
+            self._body_text)
+        self._finished_fun()
+
+
+def _launch_email_thread(host, subject, to_addr, from_addr, body_text,
+        finished_fun):
+    thread = EmailSenderThread(
+        host, subject, to_addr, from_addr, body_text, finished_fun)
+    thread.start()
+    return thread
+
+
+class FinishedThreads(object):
+
+    def __init__(self):
+        self._finished = 0
+        self._lock = threading.Lock()
+
+    def signal(self):
+        with self._lock:
+            self._finished = self._finished + 1
+            logger.info('Number of messages sent: %d.' % self._finished)
+
+
+def _send_messages(args):
+    host = args.server
+    subject = "Message from Soledad script"
+    to_addr = args.target_address
+    from_addr = args.target_address
+    body_text = "Test message"
+
+    semaphore = threading.Semaphore(args.threads)
+    threads = []
+    finished_threads = FinishedThreads()
+
+    def _finished_fun():
+        semaphore.release()
+        finished_threads.signal()
+
+    for i in xrange(args.number_of_messages):
+        semaphore.acquire()
+        threads.append(
+            _launch_email_thread(
+               host, subject, to_addr, from_addr, body_text,
+               _finished_fun))
+
+    for t in threads:
+        t.join()
+
+
+if __name__ == "__main__":
+    args = _parse_args()
+    _send_messages(args)