Merge branch 'dev' into test

.gitignore

@@ -8,3 +8,4 @@ logs/*
 conf.py
 host_rsa_key
 sessions/*
+coco.pid

coco/app.py

@@ -8,9 +8,8 @@ import time
 import threading
 import socket
 import json
-import tracemalloc
+import signal
 
-import gc
 from jms.service import AppService
 
 from .config import Config
@@ -68,8 +67,6 @@ class Coco:
         self.replay_recorder_class = None
         self.command_recorder_class = None
         self._task_handler = None
-        tracemalloc.start()
-        self.snapshot = tracemalloc.take_snapshot()
 
     @property
     def name(self):
@@ -134,16 +131,6 @@ class Coco:
     def heartbeat(self):
         _sessions = [s.to_json() for s in self.sessions]
         tasks = self.service.terminal_heartbeat(_sessions)
-        gc.collect()
-        snapshort2 = tracemalloc.take_snapshot()
-        top_stats = snapshort2.compare_to(self.snapshot, 'traceback')
-        print("[ Top 10 differences ]")
-        for stat in top_stats[:10]:
-            print(stat.count_diff)
-        # gc.collect()
-        # print(objgraph.show_most_common_types())
-        # for i in ("User", "Client", "Session", "ProxyServer", "Transport", "Channel", "InteractiveServer", "SSHInterface", "Server", "Asset"):
-        #     print("{} objects: {}".format(i, len(objgraph.by_type(i))))
         if tasks:
             self.handle_task(tasks)
         if tasks is False:
@@ -200,7 +187,10 @@ class Coco:
             if self.config['HTTPD_PORT'] != 0:
                 self.run_httpd()
 
-            self.stop_evt.wait()
+            signal.signal(signal.SIGTERM, lambda x, y: self.shutdown())
+            while self.stop_evt.wait(5):
+                print("Coco receive term signal, exit")
+                break
         except KeyboardInterrupt:
             self.stop_evt.set()
             self.shutdown()

coco/models.py

@@ -62,7 +62,11 @@ class Client:
     def send(self, b):
         if isinstance(b, str):
             b = b.encode("utf-8")
-        return self.chan.send(b)
+        try:
+            return self.chan.send(b)
+        except OSError:
+            self.close()
+            return
 
     def recv(self, size):
         return self.chan.recv(size)

coco/proxy.py

@@ -94,7 +94,7 @@ class ProxyServer:
                 timeout=TIMEOUT, compress=True, auth_timeout=10,
                 look_for_keys=False
             )
-        except (paramiko.AuthenticationException, paramiko.BadAuthenticationType, SSHException):
+        except (paramiko.AuthenticationException, paramiko.BadAuthenticationType, SSHException, TimeoutError):
             admins = self.app.config['ADMINS'] or 'administrator'
             self.client.send(warning(wr(
                 "Authenticate with server failed, contact {}".format(admins),

cocod

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+#
+
+import os
+import sys
+import argparse
+import time
+import signal
+
+from coco import Coco
+
+try:
+    from conf import config
+except ImportError:
+    print("Please prepare config file `cp conf_example.py conf.py`")
+    sys.exit(1)
+
+try:
+    os.mkdir("logs")
+    os.mkdir("keys")
+    os.mkdir("sessions")
+except:
+    pass
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+DAEMON = False
+PID_FILE = os.path.join(BASE_DIR, 'coco.pid')
+
+coco = Coco()
+coco.config.from_object(config)
+
+
+def check_pid(pid):
+    """ Check For the existence of a unix pid. """
+    try:
+        os.kill(pid, 0)
+    except OSError:
+        return False
+    else:
+        return True
+
+
+def start():
+    print("Start coco process")
+    if DAEMON:
+        fork_daemon()
+    with open(PID_FILE, 'w') as f:
+        f.write(str(os.getpid()))
+    coco.run_forever()
+
+
+def stop():
+    print("Stop coco process")
+    pid = None
+    if os.path.isfile(PID_FILE):
+        with open(PID_FILE) as f:
+            pid = f.read().strip()
+
+    if pid and pid.isdigit():
+        for i in range(15):
+            try:
+                os.kill(int(pid), signal.SIGTERM)
+            except ProcessLookupError:
+                pass
+            if check_pid(int(pid)):
+                time.sleep(1)
+                continue
+            else:
+                os.unlink(PID_FILE)
+                break
+
+
+def show_status():
+    pid = None
+    if os.path.isfile(PID_FILE):
+        with open(PID_FILE) as f:
+            pid = f.read().strip()
+
+    if pid and pid.isdigit() and check_pid(int(pid)):
+        print("Coco is running: {}".format(pid))
+    else:
+        print("Coco is stopped")
+
+
+def fork_daemon():
+    try:
+        if os.fork() > 0:
+            sys.exit(0)
+    except OSError as e:
+        sys.stderr.write("fork #1 failed: %d (%s)\n" % (e.errno, e.strerror))
+        sys.exit(1)
+
+    os.chdir(BASE_DIR)
+    os.setsid()
+    os.umask(0)
+
+    try:
+        pid = os.fork()
+        if pid > 0:
+            sys.exit(0)
+    except OSError as e:
+        sys.stderr.write("fork #2 failed: %d (%s)\n" % (e.errno, e.strerror))
+        sys.exit(1)
+
+    sys.stdout.flush()
+    sys.stderr.flush()
+    si = open('/dev/null', 'r')
+    so = open('/tmp/a.log', 'a')
+    se = open('/dev/null', 'a')
+    os.dup2(si.fileno(), sys.stdin.fileno())
+    os.dup2(so.fileno(), sys.stdout.fileno())
+    os.dup2(se.fileno(), sys.stderr.fileno())
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description="""
+            coco service control tools;
+
+            Example: \r\n 
+
+            %(prog)s start -d;
+            """
+    )
+    parser.add_argument(
+        'action', type=str, default='start',
+        choices=("start", "stop", "restart", "status"),
+        help="Action to run"
+    )
+
+    parser.add_argument('-d', '--daemon', nargs="?", const=1)
+    args = parser.parse_args()
+
+    if args.daemon:
+        DAEMON = True
+
+    action = args.action
+    if action == "start":
+        start()
+    elif action == "stop":
+        stop()
+    elif action == "restart":
+        stop()
+        DAEMON = True
+        start()
+    else:
+        show_status()

run_server.py

@@ -2,33 +2,17 @@
 # -*- coding: utf-8 -*-
 #
 
-import os
 import sys
+import subprocess
+import os
 
-from coco import Coco
-
-try:
-    from conf import config
-except ImportError:
-    print("Please prepare config file `cp conf_example.py conf.py`")
-    sys.exit(1)
-
-try:
-    os.mkdir("logs")
-    os.mkdir("keys")
-    os.mkdir("sessions")
-except:
-    pass
-
-
-coco = Coco()
-coco.config.from_object(config)
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 
-# Todo:
-# 0. argparser
-# 1. register application user
-# 2. backup record file
-# 3. xxx
 
 if __name__ == '__main__':
-    coco.run_forever()
+    try:
+        subprocess.call(['./cocod', 'start'], stdin=sys.stdin,
+                        stdout=sys.stdout, stderr=sys.stderr)
+    except KeyboardInterrupt:
+        subprocess.call(['./cocod', 'stop'], stdin=sys.stdin,
+                        stdout=sys.stdout, stderr=sys.stderr)