# Copyright (c) 2016-2017 VMware, Inc.  All rights reserved. -- VMware Confidential

import base64
import binascii
import psycopg2
from subprocess import Popen, PIPE

XML_CFG_SCRIPT = "/usr/lib/vmware-vpx/py/xmlcfg.py"
VPXD_CFG_FILE = "/etc/vmware-vpx/vpxd.cfg"
SYMKEY_FILE = "/etc/vmware-vpx/ssl/symkey.dat"

def get_odbc_ini_param(pattern):
   '''
   Get specific parameter from /etc/odbc.ini

   @pattern: The pattern to search for in odbc.ini settings.
   @return: The parameter value or None if the parameter is not found.
   '''
   result = None
   with open("/etc/odbc.ini") as f:
      for line in f:
         if pattern in line:
            result = line.split("=")[1].strip()
            break
   return result

def get_db_pass_symkey():
   '''
   Get vc database user password from vpxd.cfg using symkey.dat
   This implementation is used in Win/vCSA 6.x source.

   @return: The password in plain text.
   '''
   key3_cmd = [XML_CFG_SCRIPT, "-f", VPXD_CFG_FILE, "get",
               "/config/registry/DB/key_3"]
   p = Popen(key3_cmd, stdout=PIPE)
   f = p.communicate()[0]
   key3_val = f.strip()[1:]
   return decrypt_key3(get_symkey(), key3_val)

def decrypt_key3(symkey_hexstring, encrypted_base64):
   '''
   Decrypt password using symkey and encrypted value.

   @symkey_hexstring: symkey.dat contents
   @encrypted_base64: stripped key_3 value from vpxd.cfg
   @return: decrypted password
   '''
   from Crypto.Cipher import AES

   symkey = binascii.unhexlify(symkey_hexstring)
   encrypted_raw = base64.b64decode(encrypted_base64)
   assert (len(encrypted_raw) > 16), "Wrong encrypted string length."
   ciphertext = encrypted_raw[16:]
   iv = encrypted_raw[0:16]
   cipher = AES.new(symkey, AES.MODE_CBC, iv)
   try:
      padded_plaintext = cipher.decrypt(ciphertext)
   except Exception as e:
      print("Failed to decrypt password.")
   padding_len = ord(chr(padded_plaintext[-1]))
   plaintext = padded_plaintext[:-padding_len]
   return plaintext

def get_symkey():
   '''
   Read encryption hex string from symkey.dat

   @return: encryption key
   '''
   with open(SYMKEY_FILE) as f:
      content = f.readlines()
   content_str = content[0].strip()
   symkey_length = len(content_str)
   if symkey_length != 64:
      print("Unexpected symkey length.")
   return content_str

def getConnection(password, port):
   '''
   Connect to the local VCDB and return a cursor to it

   @return: Cursor to VCDB.
   '''
   password = password.decode('utf-8')
   connStr = "host='localhost' dbname='VCDB' user='vc' password='%s'" % password
   if port:
      connStr += " port='%s'" % port
   currentConnection = psycopg2.connect(connStr)
   currentConnection.autocommit = True
   return currentConnection.cursor()

if __name__ == '__main__':
   conn = getConnection(get_db_pass_symkey(), get_odbc_ini_param('SERVER_PORT'))

