From 77358a197d0694056debdffaa15ce96fc039f1d1 Mon Sep 17 00:00:00 2001
From: sh0ked <nikolay.ustinov90@gmail.com>
Date: Wed, 22 Mar 2017 00:28:02 +0700
Subject: [PATCH] Save screencast only for failed sessions

---
 config_template.py |  2 ++
 core/sessions.py   | 11 +++++-----
 core/video.py      | 55 ++++++++++++++++++++++++----------------------
 3 files changed, 37 insertions(+), 31 deletions(-)

diff --git a/config_template.py b/config_template.py
index ffa87de2..8eada906 100644
--- a/config_template.py
+++ b/config_template.py
@@ -64,3 +64,5 @@ class Config(object):
     VMMASTER_AGENT_PORT = 9000
 
     THREAD_POOL_MAX = 100
+    SAVE_SCREENCAST_ON_FAILED = False
+    SCREENCAST_CPU_LIMIT = 10
diff --git a/core/sessions.py b/core/sessions.py
index 0f25b8d5..db92a281 100644
--- a/core/sessions.py
+++ b/core/sessions.py
@@ -132,6 +132,10 @@ def close(self, reason=None):
         if self.vnc_helper:
             self.vnc_helper.stop_recording()
             self.vnc_helper.stop_proxy()
+            if self.take_screencast:
+                self.vnc_helper.convert_video()
+            elif getattr(config, "SAVE_SCREENCAST_ON_FAILED", False) and "failed" in self.status:
+                self.vnc_helper.convert_video()
 
         current_app.sessions.remove(self)
 
@@ -169,11 +173,8 @@ def run(self, endpoint):
         self.endpoint = endpoint
         self.set_vm(endpoint)
         self.status = "running"
-        self.vnc_helper = VNCVideoHelper(self.endpoint_ip,
-                                         filename_prefix=self.id)
-
-        if self.take_screencast:
-            self.vnc_helper.start_recording()
+        self.vnc_helper = VNCVideoHelper(self.endpoint_ip, filename_prefix=self.id)
+        self.vnc_helper.start_recording()
 
         log.info("Session %s starting on %s (%s)." %
                  (self.id, self.endpoint_name, self.endpoint_ip))
diff --git a/core/video.py b/core/video.py
index 7ec22b00..5d9538b3 100644
--- a/core/video.py
+++ b/core/video.py
@@ -1,24 +1,24 @@
-##
-##  flvrec.py - VNC to FLV recording tool.
-##
-##  Copyright (c) 2009-2010 by Yusuke Shinyama
-##
+# coding: utf-8
 
+import os
+import sys
+import signal
+import socket
 import logging
+import os.path
+import subprocess
+import websockify
 import multiprocessing
 from twisted.internet import threads
-import websockify
 
 from core.config import config
-
-import sys, socket, os, os.path, subprocess, signal
 from vnc2flv import flv, rfb, video
 from core.utils.network_utils import get_free_port
 
 log = logging.getLogger(__name__)
 
 
-class VNCVideoHelper():
+class VNCVideoHelper:
     recorder = None
     proxy = None
     __proxy_port = None
@@ -33,27 +33,29 @@ def __init__(self, host, port=5900, filename_prefix='vnc'):
         self.host = host
         self.port = port
 
-
     @staticmethod
     def _flvrec(filename, host='localhost', port=5900,
-               framerate=12, keyframe=120,
-               preferred_encoding=(0,),
-               blocksize=32, clipping=None,
-               debug=0, verbose=0):
+                framerate=12, keyframe=120,
+                preferred_encoding=(0,),
+                blocksize=32, clipping=None,
+                debug=0, verbose=0):
         fp = file(filename, 'wb')
-        pwdcache = rfb.PWDCache('%s:%d' % (host,port))
+        pwdcache = rfb.PWDCache('%s:%d' % (host, port))
         writer = flv.FLVWriter(fp, framerate=framerate, debug=debug)
-        sink = video.FLVVideoSink(writer,
-                            blocksize=blocksize, framerate=framerate, keyframe=keyframe,
-                            clipping=clipping, debug=debug)
-        client = rfb.RFBNetworkClient(host, port, sink, timeout=500/framerate,
-                                  pwdcache=pwdcache, preferred_encoding=preferred_encoding,
-                                  debug=debug)
+        sink = video.FLVVideoSink(
+            writer,
+            blocksize=blocksize, framerate=framerate, keyframe=keyframe,
+            clipping=clipping, debug=debug)
+        client = rfb.RFBNetworkClient(
+            host, port, sink, timeout=500/framerate,
+            pwdcache=pwdcache, preferred_encoding=preferred_encoding,
+            debug=debug)
         if verbose:
             log.debug('Start vnc recording to %s' % filename)
         retval = 0
         try:
             def sigint_handler(sig, frame):
+                log.debug("%s %s" % (sig, frame))
                 raise KeyboardInterrupt
             signal.signal(signal.SIGINT, sigint_handler)
             client.open()
@@ -90,7 +92,7 @@ def _flv2webm(self):
                 "/usr/bin/cpulimit",
                 "-z",
                 "-b",
-                "-l", "15",
+                "-l", getattr(config, "SCREENCAST_CPU_LIMIT", 10),
                 "-p", "%s" % converter.pid
             ], stdin=subprocess.PIPE)
 
@@ -99,7 +101,7 @@ def _flv2webm(self):
             self.delete_source_video()
 
     def delete_source_video(self):
-        if os.path.isfile('%s.webm' % self.__filepath.split('.flv')[0]):
+        if self.__filepath and os.path.isfile(self.__filepath):
             os.remove(self.__filepath)
             log.debug('Source video %s was deleted' % self.__filepath)
 
@@ -146,9 +148,10 @@ def start_recording(self, framerate=5, size=(800, 600)):
                                                 kwargs=kwargs)
         self.recorder.start()
 
+    def convert_video(self):
+        d = threads.deferToThread(self._flv2webm)
+        d.addBoth(lambda s: None)
+
     def stop_recording(self):
         if self.recorder and self.recorder.is_alive():
             self.recorder.terminate()
-
-            d = threads.deferToThread(self._flv2webm)
-            d.addBoth(lambda s: None)
\ No newline at end of file