Initial commit (Vesion 0.1)

[tablet-suite] / src / pcsutils.py

diff --git a/src/pcsutils.py b/src/pcsutils.py
new file mode 100755 (executable)
index 0000000..7d8a930
--- /dev/null
+++ b/src/pcsutils.py
@@ -0,0 +1,165 @@
+import os.path
+
+import commands
+import os
+import pwd
+import settings
+import socket
+import sys
+
+import paramiko
+
+from backup.pcsbackuputils import createFolder
+
+sshPath = os.path.expanduser('~/.ssh/') 
+known_hosts = os.path.join(sshPath, 'known_hosts')
+log_file = os.path.expanduser('~/.pcsuite/.ssh_log')
+user = 'root'
+keyName = 'rsa_key'
+
+def create_route(host, port=22):
+    # Verify Auth with privateKey
+    try:
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.settimeout(15)
+        sock.connect((host, port))
+        sock.close()
+        return True
+    except:
+        print 'No route to host'
+        return False
+
+def verify_exist_keys(host, port=22):
+    try:    
+        transport = _create_transport(host, port)
+    except:
+        return False
+    try:
+        getKey = paramiko.RSAKey.from_private_key_file(sshPath + keyName)
+        transport.start_client()
+        transport.auth_publickey(user, getKey)
+        if transport.is_authenticated():
+            transport.close()
+            return True
+    except:
+        # 'Error in auth with publickey, try with password...'
+        return False
+    return False
+
+def keyExchange(host, passwd, port=22):
+    if not os.path.exists(sshPath):
+        createFolder(sshPath)
+
+    # Clean cached keys in ssh-agent
+    os.system('ssh-add -d')    
+
+    try:
+        transport = _create_transport(host, port)
+    except:
+        transport.close()
+        return False
+
+    if not _add_host_fingerprint(host):
+        transport.close()
+        return False
+
+    if not _authenticate(user, passwd, transport):
+        transport.close()
+        return False
+    
+    if not _add_key_to_host(host, transport):
+        transport.close()
+        return False
+
+    transport.stop_thread()
+    transport.close()
+    return True
+
+def initDirs():
+    settings.makeDirs()
+
+def _create_transport(host, port):
+    # Create a transport and initiate client mode
+    try:
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.settimeout(15)
+        sock.connect((host, port))
+    except Exception, msg:
+        print 'Connect failed: ' + str(msg)
+        raise Exception('Error while create sockets.')
+    transport = paramiko.Transport(sock)
+    return transport
+
+def _add_host_fingerprint(host):
+    if not os.path.exists(known_hosts):
+        os.system('touch %s' %known_hosts)
+    if os.system('ssh-keyscan -t rsa %s >> %s' %(host, known_hosts)) != 0:
+        return False
+    return True
+
+def _generate_keys():
+    # Generate public and private RSAKey
+    keyFile = os.path.join(sshPath, keyName)
+    if not os.path.exists(keyFile):
+        privateKey = paramiko.RSAKey.generate(2048)
+        privateKey.write_private_key_file(keyFile)
+        login = pwd.getpwuid(os.geteuid())[0]
+        publicKey = '%s %s %s@%s' %(privateKey.get_name(), 
+                                    privateKey.get_base64(),
+                                    login , socket.gethostname())
+        try:
+            keyFile = open(keyFile + '.pub','w')
+            keyFile.write(publicKey)
+            keyFile.close()
+        except:
+            print 'Error while save the public key'
+            raise Exception()
+    else:
+        try:
+            privateKey = paramiko.RSAKey.from_private_key_file(keyFile)
+            login = pwd.getpwuid(os.geteuid())[0]
+            publicKey = '%s %s %s@%s' %(privateKey.get_name(), 
+                                        privateKey.get_base64(),
+                                    login , socket.gethostname())
+        except:
+            print 'Error while read the private key'
+            raise Exception()
+    return publicKey
+
+def _authenticate(user, passwd, transport):
+    # Try Auth with password
+    try:
+        transport.start_client()
+        transport.auth_password(user, passwd)
+    except:
+        print 'Verify user or password.'
+        return False
+    if not transport.is_authenticated():
+        print 'Authentication fail'
+        return False
+
+    try:
+        exception = transport.get_exception()
+        if exception:
+            raise exception
+    except Exception, msg:
+        print 'Error in connection: ' + str(msg)
+        return False
+    return True
+
+def _add_key_to_host(host, transport):
+    # Add publickey in host
+    if not transport.is_active():
+        print 'Channel is not active'
+        return False
+    
+    paramiko.util.log_to_file(log_file, 10)
+    channel = transport.open_session()
+    try:
+        channel.exec_command('mkdir -p ~/.ssh; echo %s >> .ssh/authorized_keys' % (_generate_keys()))
+    except Exception, msg:
+        print 'Error while generate or add the keys.'
+        channel.close()
+        return False
+    channel.close()
+    return True