From c1437353934245851de2b51f823cf8951287fd47 Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Wed, 31 Aug 2016 01:00:20 +0800
Subject: [PATCH] Add ssh-key-gen function

---
 apps/jumpserver/settings.py               |  9 +++--
 apps/users/models.py                      |  8 ++--
 apps/users/templates/users/user_list.html |  2 +-
 apps/users/utils.py                       | 49 +++++++++++++++++++++++
 requirements.txt                          |  1 +
 5 files changed, 61 insertions(+), 8 deletions(-)

diff --git a/apps/jumpserver/settings.py b/apps/jumpserver/settings.py
index 6d6eae6c6..7dfc1e603 100644
--- a/apps/jumpserver/settings.py
+++ b/apps/jumpserver/settings.py
@@ -96,7 +96,7 @@ TEMPLATES = [
     },
 ]
 
-WSGI_APPLICATION = 'jumpserver.wsgi.application'
+# WSGI_APPLICATION = 'jumpserver.wsgi.application'
 
 # Database
 # https://docs.djangoproject.com/en/1.10/ref/settings/#databases
@@ -230,9 +230,6 @@ MEDIA_URL = '/media/'
 
 MEDIA_ROOT = os.path.join(BASE_DIR, 'media').replace('\\', '/') + '/'
 
-# Custom User Auth model
-AUTH_USER_MODEL = 'users.User'
-
 # Use django-bootstrap-form to format template, input max width arg
 BOOTSTRAP_COLUMN_COUNT = 11
 
@@ -275,3 +272,7 @@ WS4REDIS_PREFIX = 'demo'
 SESSION_ENGINE = 'redis_sessions.session'
 
 SESSION_REDIS_PREFIX = 'session'
+
+# Custom User Auth model
+AUTH_USER_MODEL = 'users.User'
+
diff --git a/apps/users/models.py b/apps/users/models.py
index b086780d2..2681a2470 100644
--- a/apps/users/models.py
+++ b/apps/users/models.py
@@ -10,9 +10,8 @@ from django.db import models
 from django.contrib.auth.models import AbstractUser, Permission
 from django.db.models.signals import post_save
 from django.dispatch import receiver
-from rest_framework.authtoken.models import Token
 from django.db import IntegrityError
-
+from rest_framework.authtoken.models import Token
 
 
 # class Role(models.Model):
@@ -250,5 +249,8 @@ def generate_fake():
 @receiver(post_save, sender=settings.AUTH_USER_MODEL)
 def create_auth_token(sender, instance=None, created=False, **kwargs):
     if created:
-        Token.objects.create(user=instance)
+        try:
+            Token.objects.create(user=instance)
+        except IntegrityError:
+            pass
 
diff --git a/apps/users/templates/users/user_list.html b/apps/users/templates/users/user_list.html
index 1c61d3cdf..e1c8da5cd 100644
--- a/apps/users/templates/users/user_list.html
+++ b/apps/users/templates/users/user_list.html
@@ -42,7 +42,7 @@
              </td>
              <td class="text-center">
                  <a href="{% url 'users:user-edit' pk=user.id %}" class="btn btn-xs btn-info">编辑</a>
-                 <a href="{% url 'users:user-delete' pk=user.id %}" class="btn btn-xs btn-danger del {% if user.username == 'admin' %} disabled {% endif %}">删除</a>
+                 <a href="{% url 'users:user-delete' pk=user.id %}" class="btn btn-xs btn-danger del {% if user.id == request.user.id or user.username == 'admin' %} disabled {% endif %}">删除</a>
              </td>
          </tr>
      {% endfor %}
diff --git a/apps/users/utils.py b/apps/users/utils.py
index 8ed37c599..d6cee1080 100644
--- a/apps/users/utils.py
+++ b/apps/users/utils.py
@@ -1,12 +1,61 @@
 # ~*~ coding: utf-8 ~*~
 #
 
+import os
+import logging
+
+from paramiko.rsakey import RSAKey
 from django.contrib.auth.mixins import UserPassesTestMixin
 from django.urls import reverse_lazy
 
+try:
+    import cStringIO as StringIO
+except ImportError:
+    import StringIO
+
+
+logger = logging.getLogger('jumpserver')
+
 
 class AdminUserRequiredMixin(UserPassesTestMixin):
     login_url = reverse_lazy('users:login')
 
     def test_func(self):
         return self.request.user.is_staff
+
+
+def ssh_key_gen(length=2048, password=None, username='root', hostname=None):
+    """Generate user ssh private and public key
+
+    Use paramiko RSAKey generate it.
+
+    """
+
+    if hostname is None:
+        hostname = os.uname()[1]
+
+    f = StringIO.StringIO()
+
+    try:
+        logger.debug('Begin to generate ssh private key ...')
+        private_key_obj = RSAKey.generate(length)
+        private_key_obj.write_private_key(f, password=password)
+        private_key = f.getvalue()
+
+        public_key = "%(key_type)s %(key_content)s %(username)s@%(hostname)s" % {
+            'key_type': private_key_obj.get_name(),
+            'key_content': private_key_obj.get_base64(),
+            'username': username,
+            'hostname': hostname,
+        }
+
+        logger.debug('Finish to generate ssh private key ...')
+        return private_key, public_key
+
+    except IOError:
+        raise IOError('These is error when generate ssh key.')
+
+
+
+
+
diff --git a/requirements.txt b/requirements.txt
index 34c61c4ce..587da6a34 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -13,3 +13,4 @@ wcwidth==0.1.7
 websocket-client==0.37.0
 djangorestframework==3.4.5
 ForgeryPy==0.1
+paramiko==2.0.2