From 4f05b88abf8f005735a7d799ba06a007e2f9a6ea Mon Sep 17 00:00:00 2001
From: Jake Howard <git@theorangeone.net>
Date: Wed, 12 Dec 2018 15:53:56 +0000
Subject: [PATCH] Init CLI

---
 catfish/__init__.py        |  3 +++
 catfish/__main__.py        | 39 ++++++++++++++++++++++++++++++++++++++
 catfish/utils/processes.py | 15 +++++++++++++++
 catfish/worker/__init__.py | 22 +++++++++++++++++++++
 scripts/test.sh            |  2 +-
 setup.py                   |  5 +++++
 tests/__init__.py          |  9 +++++++++
 tests/test_main.py         | 10 ++++++++++
 8 files changed, 104 insertions(+), 1 deletion(-)
 create mode 100644 catfish/__main__.py
 create mode 100644 catfish/utils/processes.py
 create mode 100644 catfish/worker/__init__.py
 create mode 100644 tests/test_main.py

diff --git a/catfish/__init__.py b/catfish/__init__.py
index e69de29..14bd544 100644
--- a/catfish/__init__.py
+++ b/catfish/__init__.py
@@ -0,0 +1,3 @@
+import pkg_resources
+
+__version__ = pkg_resources.require("catfish")[0].version
diff --git a/catfish/__main__.py b/catfish/__main__.py
new file mode 100644
index 0000000..acfc440
--- /dev/null
+++ b/catfish/__main__.py
@@ -0,0 +1,39 @@
+import click
+import daemonize
+
+from catfish import __version__, worker
+from catfish.utils.processes import terminate_processes
+
+
+@click.group()
+@click.version_option(__version__, prog_name="catfish")
+def cli():
+    pass
+
+
+@cli.command()
+@click.option("--no-fork", is_flag=True)
+def start(no_fork):
+    daemon = daemonize.Daemonize(
+        "catfish", worker.PID_FILE, worker.main, foreground=no_fork
+    )
+    try:
+        daemon.start()
+    except SystemExit:
+        return
+    proc = worker.get_running_process()
+    click.echo("Worker started with pid {}".format(proc.pid))
+
+
+@cli.command()
+def stop():
+    if not worker.is_running():
+        click.echo("Worker not running")
+        return
+    proc = worker.get_running_process()
+    click.echo("Terminating process {}".format(proc.pid))
+    terminate_processes([proc])
+
+
+if __name__ == "__main__":
+    cli()
diff --git a/catfish/utils/processes.py b/catfish/utils/processes.py
new file mode 100644
index 0000000..da9ef10
--- /dev/null
+++ b/catfish/utils/processes.py
@@ -0,0 +1,15 @@
+from typing import List
+
+import psutil
+
+
+def terminate_processes(procs: List[psutil.Process], timeout=3):
+    for p in procs:
+        p.terminate()
+    gone, alive = psutil.wait_procs(procs, timeout=timeout)
+
+    if alive:
+        # send SIGKILL
+        for p in alive:
+            p.kill()
+        gone, alive = psutil.wait_procs(alive, timeout=timeout)
diff --git a/catfish/worker/__init__.py b/catfish/worker/__init__.py
new file mode 100644
index 0000000..a79df02
--- /dev/null
+++ b/catfish/worker/__init__.py
@@ -0,0 +1,22 @@
+import time
+import os
+import tempfile
+import psutil
+
+
+PID_FILE = os.path.join(tempfile.gettempdir(), "catfish.pid")
+
+
+def is_running():
+    return os.path.exists(PID_FILE)
+
+
+def get_running_process() -> psutil.Process:
+    assert is_running()
+    with open(PID_FILE) as f:
+        return psutil.Process(int(f.read()))
+
+
+def main():
+    while True:
+        time.sleep(1)
diff --git a/scripts/test.sh b/scripts/test.sh
index 0c12b1c..168fd48 100755
--- a/scripts/test.sh
+++ b/scripts/test.sh
@@ -9,4 +9,4 @@ nose2 $@ -C --coverage catfish --verbose
 black catfish tests setup.py --check
 flake8 catfish tests setup.py --ignore=E128,E501
 isort -rc -c catfish tests setup.py
-mypy --strict-optional catfish
+mypy --strict-optional --ignore-missing-imports catfish
diff --git a/setup.py b/setup.py
index b60cc50..fdc08ef 100644
--- a/setup.py
+++ b/setup.py
@@ -11,5 +11,10 @@ setup(
     include_package_data=True,
     zip_safe=False,
     pathon_requires=">=3.6",
+    install_requires=["click", "daemonize", "psutil"],
+    entry_points="""
+        [console_scripts]
+        ctf=catfish.__main__:cli
+    """,
     project_urls={"GitHub: Issues": "https://github.com/realorangeone/catfish/issues"},
 )
diff --git a/tests/__init__.py b/tests/__init__.py
index e69de29..6c83a8d 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -0,0 +1,9 @@
+from unittest import TestCase
+from click.testing import CliRunner
+from catfish.__main__ import cli
+
+
+class BaseTestCase(TestCase):
+    def setUp(self):
+        self.cli_runner = CliRunner()
+        self.cli = cli
diff --git a/tests/test_main.py b/tests/test_main.py
new file mode 100644
index 0000000..9173b0c
--- /dev/null
+++ b/tests/test_main.py
@@ -0,0 +1,10 @@
+from catfish import __version__
+from tests import BaseTestCase
+
+
+class MainCLITestCase(BaseTestCase):
+    def test_version(self):
+        result = self.cli_runner.invoke(self.cli, ["--version"])
+        self.assertEqual(result.exit_code, 0)
+        self.assertIn(__version__, result.output)
+        self.assertIn("catfish", result.output)