From 73b5955d604788f1d6db041cdf53e3c07a3376b4 Mon Sep 17 00:00:00 2001
From: Andrew Kay <git@ajk.me>
Date: Sun, 1 Dec 2019 15:59:49 -0600
Subject: [PATCH] Add SIGINT and SIGTERM signal handler

---
 oec/__main__.py   | 18 ++++++++++++++++++
 oec/controller.py | 21 +++++++++++++++++----
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/oec/__main__.py b/oec/__main__.py
index ff6c128..3aeba4c 100644
--- a/oec/__main__.py
+++ b/oec/__main__.py
@@ -1,4 +1,5 @@
 import time
+import signal
 import logging
 import argparse
 from serial import Serial
@@ -12,6 +13,8 @@ from .keymap_3483 import KEYMAP as KEYMAP_3483
 
 logging.basicConfig(level=logging.INFO)
 
+controller = None
+
 def _get_keymap(terminal_id, extended_id):
     keymap = KEYMAP_3278_2
 
@@ -31,7 +34,22 @@ def _create_session(args, terminal):
 
     raise ValueError('Unsupported emulator')
 
+def _signal_handler(number, frame):
+    global controller
+
+    print('Stopping controller...')
+
+    if controller:
+        controller.stop()
+
+        controller = None
+
+signal.signal(signal.SIGINT, _signal_handler)
+signal.signal(signal.SIGTERM, _signal_handler)
+
 def main():
+    global controller
+
     parser = argparse.ArgumentParser(description=('An open replacement for the IBM 3174 '
                                                   'Establishment Controller'))
 
diff --git a/oec/controller.py b/oec/controller.py
index b7e0190..ec403a7 100644
--- a/oec/controller.py
+++ b/oec/controller.py
@@ -69,6 +69,14 @@ class Controller:
             if poll_response:
                 self._handle_poll_response(poll_response)
 
+        self._terminate_session()
+
+        if self.terminal:
+            self.terminal = None
+
+    def stop(self):
+        self.running = False
+
     def _handle_terminal_attached(self, poll_response):
         self.logger.info('Terminal attached')
 
@@ -99,10 +107,7 @@ class Controller:
     def _handle_terminal_detached(self):
         self.logger.info('Terminal detached')
 
-        if self.session:
-            self.session.terminate()
-
-            self.session = None
+        self._terminate_session()
 
         self.terminal = None
 
@@ -119,6 +124,14 @@ class Controller:
 
         self.session.start()
 
+    def _terminate_session(self):
+        if not self.session:
+            return
+
+        self.session.terminate()
+
+        self.session = None
+
     def _handle_poll_response(self, poll_response):
         if isinstance(poll_response, KeystrokePollResponse):
             self._handle_keystroke_poll_response(poll_response)