This commit is contained in:
Jake Howard 2018-12-12 15:53:56 +00:00
parent 198eb1f209
commit 4f05b88abf
Signed by: jake
GPG key ID: 57AFB45680EDD477
8 changed files with 104 additions and 1 deletions

View file

@ -0,0 +1,3 @@
import pkg_resources
__version__ = pkg_resources.require("catfish")[0].version

39
catfish/__main__.py Normal file
View file

@ -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()

View file

@ -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)

View file

@ -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)

View file

@ -9,4 +9,4 @@ nose2 $@ -C --coverage catfish --verbose
black catfish tests setup.py --check black catfish tests setup.py --check
flake8 catfish tests setup.py --ignore=E128,E501 flake8 catfish tests setup.py --ignore=E128,E501
isort -rc -c catfish tests setup.py isort -rc -c catfish tests setup.py
mypy --strict-optional catfish mypy --strict-optional --ignore-missing-imports catfish

View file

@ -11,5 +11,10 @@ setup(
include_package_data=True, include_package_data=True,
zip_safe=False, zip_safe=False,
pathon_requires=">=3.6", 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"}, project_urls={"GitHub: Issues": "https://github.com/realorangeone/catfish/issues"},
) )

View file

@ -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

10
tests/test_main.py Normal file
View file

@ -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)