Add management commands

Author: skatsaounis

patchman/apps.py

@@ -0,0 +1,7 @@
+from __future__ import unicode_literals
+
+from django.apps import AppConfig
+
+
+class PatchmanConfig(AppConfig):
+    name = 'patchman'

patchman/management/__init__.py

@@ -0,0 +1,33 @@
+from django.contrib.auth.management.commands import createsuperuser
+from django.core.management import CommandError
+
+
+class Command(createsuperuser.Command):
+    help = 'Crate a superuser, and allow password to be provided'
+
+    def add_arguments(self, parser):
+        super(Command, self).add_arguments(parser)
+        parser.add_argument(
+            '--password', dest='password', default=None,
+            help='Specifies the password for the superuser.',
+        )
+
+    def handle(self, *args, **options):
+        password = options.get('password')
+        username = options.get('username')
+        database = options.get('database')
+
+        if options['interactive']:
+            raise CommandError(
+                'Command is required to run with --no-input option')
+        if password and not username:
+            raise CommandError(
+                '--username is required if specifying --password')
+
+        super(Command, self).handle(*args, **options)
+
+        if password:
+            user = self.UserModel._default_manager.db_manager(
+                database).get(username=username)
+            user.set_password(password)
+            user.save()

patchman/management/commands/__init__.py

@@ -0,0 +1,17 @@
+from django.core.management.base import BaseCommand, CommandError
+from hosts.models import Host
+
+
+class Command(BaseCommand):
+    help = 'Enable/Disable rDNS check for hosts'
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            '--disable', action='store_false', default=True, dest='rdns_check',
+            help='If set, disables rDNS check')
+
+    def handle(self, *args, **options):
+        try:
+            Host.objects.all().update(check_dns=options['rdns_check'])
+        except Exception as e:
+            raise CommandError('Failed to update rDNS check', str(e))

patchman/management/commands/createsuperuser_with_password.py

@@ -0,0 +1,52 @@
+import os
+import re
+import sys
+import codecs
+from random import choice
+from tempfile import NamedTemporaryFile
+from shutil import copy
+
+from django.core.management.base import BaseCommand
+
+
+class Command(BaseCommand):
+    help = 'Set SECRET_KEY of Patchman Application.'
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            '--key', help=(
+                'The SECRET_KEY to be used by Patchman. If not set, a random '
+                'key of length 50 will be created.'))
+
+    @staticmethod
+    def get_random_key():
+        chars = 'abcdefghijklmnopqrstuvwxyz0123456789!@#$%^&*(-_=+)'
+        return ''.join([choice(chars) for i in range(50)])
+
+    def handle(self, *args, **options):
+        secret_key = options.get('key', self.get_random_key())
+
+        if sys.prefix == '/usr':
+            conf_path = '/etc/patchman'
+        else:
+            conf_path = os.path.join(sys.prefix, 'etc/patchman')
+            # if conf_path doesn't exist, try ./etc/patchman
+            if not os.path.isdir(conf_path):
+                conf_path = './etc/patchman'
+        local_settings = os.path.join(conf_path, 'local_settings.py')
+
+        settings_contents = codecs.open(
+            local_settings, 'r', encoding='utf-8').read()
+        settings_contents = re.sub(
+            r"(?<=SECRET_KEY = ')'", secret_key + "'", settings_contents)
+
+        f = NamedTemporaryFile(delete=False)
+        temp = f.name
+        f.close()
+
+        fh = codecs.open(temp, 'w+b', encoding='utf-8')
+        fh.write(settings_contents)
+        fh.close()
+
+        copy(temp, local_settings)
+        os.remove(temp)

patchman/management/commands/set_rdns_check.py

@@ -89,6 +89,7 @@
     'repos.apps.ReposConfig',
     'reports.apps.ReportsConfig',
     'util.apps.UtilConfig',
+    'patchman.apps.PatchmanConfig',
 ]
 
 REST_FRAMEWORK = {