Stormer 0.1:

1, change project name, LocustWrapper to Stormer;
2, add sput subcommands, support copying local file/directory to remote host with ssh.

.gitignore

@@ -87,3 +87,5 @@ ENV/
 
 # Rope project settings
 .ropeproject
+
+.DS_Store

README.md

@@ -1,22 +1,40 @@
-# LocustWrapper
+# Stormer
 
-A wrapper for making locustio more convenient to use.
+Wrappers for making load test more convienient.
 
 ## install
 
-LocustWrapper is based on [`locustio`](https://github.com/locustio/locust), so you need to install `locust` dependencies first.
+Stormer is based on [`locustio`](https://github.com/locustio/locust), so you need to install `locust` dependencies first.
 
 ```bash
 $ pip install -r requirements.txt --upgrade
 ```
 
-## usage
+## usages
+
+Currently, Stormer supports two subcommands.
 
 ```text
 $ python main.py -h
-usage: main.py [-h] [-f LOCUSTFILE] [--slaves_num SLAVES_NUM]
+usage: main.py [-h] {locust,sput} ...
+
+Wrappers for making load test more convienient.
 
-A wrapper for making locustio more convienient to use.
+positional arguments:
+  {locust,sput}  sub-command help
+    locust       locust wrapper.
+    sput         scp wrapper for putting files.
+
+optional arguments:
+  -h, --help     show this help message and exit
+```
+
+`locust` usage.
+
+```text
+$ usage: main.py locust [-h] [-f LOCUSTFILE] [-P PORT] [--slaves_num SLAVES_NUM]
+
+Start locust master and specified number of slaves with one command.
 
 optional arguments:
   -h, --help            show this help message and exit
@@ -28,10 +46,30 @@ optional arguments:
                         Specify number of locust slaves.
 ```
 
+`sput` usage.
+
+```text
+$ python main.py sput -h
+usage: main.py sput [-h] [--hostsfile HOSTSFILE] [--localpath LOCALPATH] [--remotepath REMOTEPATH]
+
+Copy local file/directory to remote host with ssh.
+
+optional arguments:
+  -h, --help            show this help message and exit
+  --hostsfile HOSTSFILE
+                        Specify hosts file to handle.
+  --localpath LOCALPATH
+                        Specify localpath of file or directory to transfer.
+  --remotepath REMOTEPATH
+                        Specify remotepath of file or directory to transfer.
+```
+
 ## example
 
+Start locust master and 4 locust slaves.
+
 ```text
-$ python main.py -f examples/demo_task.py --slaves_num 4
+$ python main.py locust -f examples/demo_task.py --slaves_num 4
 [2017-02-26 10:52:04,875] Leos-MacBook-Air.local/INFO/logger: Starting Locust 0.8a2
 [2017-02-26 10:52:04,897] Leos-MacBook-Air.local/INFO/logger: Starting web monitor at *:8089
 [2017-02-26 01:32:15,757] Leos-MacBook-Air.local/INFO/locust.runners: Client 'Leos-MacBook-Air.local_9cfcb5acf942af4b52063c138952a999' reported as ready. Current
@@ -43,3 +81,10 @@ ly 3 clients ready to swarm.
 [2017-02-26 01:32:15,782] Leos-MacBook-Air.local/INFO/locust.runners: Client 'Leos-MacBook-Air.local_cc9d414341823d0e9421679b5f9dd4c4' reported as ready. Current
 ly 4 clients ready to swarm.
 ```
+
+Copy local directory to all remote hosts.
+
+```text
+$ python main.py sput --hostsfile examples/hosts.yml --localpath ~/MyProjects/test_dir --remotepa
+th /root/test_dir
+```

examples/hosts.yml

@@ -0,0 +1,15 @@
+-
+    hostname: '192.168.10.11'
+    port: 22
+    username: "root"
+    password: "123456"
+-
+    hostname: '192.168.10.12'
+    port: 22
+    username: "root"
+    password: "123456"
+-
+    hostname: '192.168.10.13'
+    port: 22
+    username: "root"
+    password: "123456"

main.py

@@ -1,27 +1,45 @@
 import sys
+import os
 import argparse
-from locust_wrapper.logger import logger
-from locust_wrapper.start import LocustStarter
+from stormer.logger import logger
+from stormer.locust import LocustStarter
+from stormer import ssh
 
-def parse_args():
-    """ parse command line options.
+def main():
+    """ parse command line options and run commands.
     """
     parser = argparse.ArgumentParser(
-        description='A wrapper for making locustio more convienient to use.')
+        description='Wrappers for making load test more convienient.')
+
+    subparsers = parser.add_subparsers(help='sub-command help')
+
+    locust_subparser = subparsers.add_parser('locust', help='locust wrapper.',
+        description='Start locust master and specified number of slaves with one command.')
+    locust_subparser.add_argument('-f', '--locustfile', dest="locustfile",
+                                  help="Specify locust file to run test.")
+    locust_subparser.add_argument('-P', '--port', '--web-port', dest="port", default='',
+                                  help="Port on which to run web host.")
+    locust_subparser.add_argument('--slaves_num', dest="slaves_num", default='1',
+                                  help="Specify number of locust slaves.")
+    locust_subparser.set_defaults(func=main_locust)
 
-    parser.add_argument('-f', '--locustfile', dest="locustfile",
-                        help="Specify locust file to run test.")
-    parser.add_argument('-P', '--port', '--web-port', dest="port", default='',
-                        help="Port on which to run web host.")
-    parser.add_argument('--slaves_num', dest="slaves_num", default='1',
-                        help="Specify number of locust slaves.")
+    sput_subparser = subparsers.add_parser('sput', help='scp wrapper for putting files.',
+        description='Copy local file/directory to remote host with ssh.')
+    sput_subparser.add_argument('--hostsfile', dest="hostsfile",
+                                help="Specify hosts file to handle.")
+    sput_subparser.add_argument('--localpath', dest="localpath",
+                                help="Specify localpath of file or directory to transfer.")
+    sput_subparser.add_argument('--remotepath', dest="remotepath",
+                                help="Specify remotepath of file or directory to transfer.")
+    sput_subparser.set_defaults(func=main_sput)
 
     args = parser.parse_args()
+    args.func(args)
+
     return args
 
-def main():
-    args = parse_args()
-    locustfile = args.locustfile.strip()
+def main_locust(args):
+    locustfile = args.locustfile
     if locustfile is None:
         logger.error("locustfile must be specified! use the -f option.")
         sys.exit(0)
@@ -30,7 +48,21 @@ def main():
     port = int(port) if port.isdigit() else None
     slaves_num = int(args.slaves_num.strip())
 
-    LocustStarter().start(locustfile=locustfile, port=port, slaves_num=slaves_num)
+    LocustStarter().start(locustfile=locustfile.strip(), port=port, slaves_num=slaves_num)
+
+def main_sput(args):
+    hostsfile = args.hostsfile
+    if hostsfile is None:
+        logger.error("hostsfile must be specified! use the --hostsfile option.")
+        sys.exit(0)
+
+    localpath = args.localpath
+    if localpath is None:
+        logger.error("localpath must be specified! use the --localpath option.")
+        sys.exit(0)
+
+    remotepath = args.remotepath.strip()
+    ssh.sput(hostsfile.strip(), localpath.strip(), remotepath)
 
 
 if __name__ == '__main__':

requirements.txt

@@ -1,17 +1,24 @@
 appdirs==1.4.0
+cffi==1.9.1
 click==6.7
+cryptography==1.7.2
 Flask==0.12
 gevent==1.2.1
 Glances==2.8.2
 greenlet==0.4.12
+idna==2.4
 itsdangerous==0.24
 Jinja2==2.9.5
 -e git+https://github.com/locustio/locust.git@master#egg=locustio
 MarkupSafe==0.23
 msgpack-python==0.4.8
 packaging==16.8
+paramiko==2.1.2
 psutil==5.1.3
+pyasn1==0.2.3
+pycparser==2.17
 pyparsing==2.1.10
+PyYAML==3.12
 pyzmq==15.2.0
 requests==2.13.0
 requests-toolbelt==0.7.1

stormer/__init__.py

@@ -0,0 +1 @@
+import gevent.monkey; gevent.monkey.patch_all()

locust_wrapper/dummy_options.py → stormer/base.py

@@ -1,3 +1,11 @@
+import yaml
+from stormer.logger import loglevel, logfile
+
+def parse_yaml(yamlfile):
+    with open(yamlfile) as stream:
+        return yaml.load(stream)
+
+
 class DummyOptions(object):
     def __init__(self):
         self.host = None
@@ -14,8 +22,8 @@ def __init__(self):
         self.num_clients = 1
         self.hatch_rate = 1
         self.num_requests = None
-        self.loglevel = "INFO"
-        self.logfile = None
+        self.loglevel = loglevel
+        self.logfile = logfile
         self.print_stats = False
         self.only_summary = False
         self.no_reset_stats = False

locust_wrapper/start.py → stormer/locust.py

@@ -7,9 +7,8 @@
 from locust import events, web
 from locust.main import version, load_locustfile
 from locust.stats import print_percentile_stats, print_error_report, print_stats
-from locust_wrapper.dummy_options import master_options, slave_options
-from locust_wrapper.logger import logger
-
+from stormer.base import master_options, slave_options
+from stormer.logger import logger
 
 def parse_locustfile(locustfile):
     docstring, locusts = load_locustfile(locustfile)

stormer/logger.py

@@ -0,0 +1,9 @@
+import logging
+from locust.log import setup_logging
+
+# setup logging
+loglevel = 'INFO'
+logfile = None
+
+setup_logging(loglevel, logfile)
+logger = logging.getLogger(__name__)

stormer/ssh.py

@@ -0,0 +1,91 @@
+import os
+import time
+import paramiko
+import threading
+from stormer.logger import logger
+from stormer import base
+
+class SSHConnector(paramiko.SSHClient):
+
+    def __init__(self, host):
+        super(SSHConnector, self).__init__()
+        self.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        self.connect(**host)
+        self.sftp = self.open_sftp()
+        logger.info("Connected to {}:{} with {}".format(
+            host['hostname'], host['port'], host['username']))
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        self.sftp.close()
+        self.close()
+
+    def exec_cmd(self, cmd):
+        stdin, stdout, stderr = self.exec_command(cmd)
+        return stdout.readlines()
+
+    def create_remote_dir(self, remote_dir):
+        cmd_create_dir = 'mkdir -p {}'.format(remote_dir)
+        self.exec_cmd(cmd_create_dir)
+        time.sleep(0.5)
+
+    def get_file(self, remotefile, localfile):
+        self.sftp.get(remotefile, localfile)
+
+    def get_dir(self, remote_dir, local_dir):
+        pass
+
+    def put_file(self, localfile, remotefile):
+        remote_dirpath = os.path.dirname(remotefile)
+        if remote_dirpath == '':
+            remotefile = os.path.join('~', remotefile)
+        else:
+            cmd_check_dir = 'test -d {} && echo "ISDIR" || echo "NOTDIR"'.format(remote_dirpath)
+            output = self.exec_cmd(cmd_check_dir)
+            if output and 'NOTDIR' in output[0]:
+                self.create_remote_dir(remote_dirpath)
+
+        logger.info('transfer {} to {}'.format(localfile, remotefile))
+        self.sftp.put(localfile, remotefile)
+
+    def __get_all_files_in_local_dir(self, local_dir):
+        all_files = list()
+        files = os.listdir(local_dir)
+        for filename in files:
+            path = os.path.join(local_dir, filename)
+            if os.path.isdir(path):
+                all_files.extend(self.__get_all_files_in_local_dir(path))
+            else:
+                all_files.append(path)
+        return all_files
+
+    def put_dir(self, local_dir, remote_dir):
+        """ transfer all file in local_dir to remote_dir.
+        """
+        all_files = self.__get_all_files_in_local_dir(local_dir)
+        for localfile in all_files:
+            relative_path = localfile[len(local_dir) + 1:]
+            remotefile = os.path.join(remote_dir, relative_path)
+            remote_file_dir = os.path.dirname(remotefile)
+            self.create_remote_dir(remote_file_dir)
+            self.put_file(localfile, remotefile)
+
+
+def sput(hostsfile, localpath, remotepath):
+
+    def _put_file(host, localfile, remotefile):
+        with SSHConnector(host) as ssh_client:
+            ssh_client.put_file(localfile, remotefile)
+
+    def _put_dir(host, local_dir, remote_dir):
+        with SSHConnector(host) as ssh_client:
+            ssh_client.put_dir(local_dir, remote_dir)
+
+    host_list = base.parse_yaml(hostsfile)
+    localpath = os.path.abspath(localpath)
+    target_func = _put_dir if os.path.isdir(localpath) else _put_file
+
+    for host in host_list:
+        threading.Thread(target=target_func, args=(host, localpath, remotepath)).start()