#!/usr/bin/env python3
# Copyright (c) 2009, 2010 Arista Networks, Inc.  All rights reserved.
# Arista Networks, Inc. Confidential and Proprietary.

# This utility is used to set the root password. We do this here, so that
# root password setting doesn't depend on Aaa/Sysdb being initialzed correctly.
# Although we do depend on startup-config not being messed up. By default root
# login is disabled before this script runs, so there is no race where root login
# can be accessed before this script can be executed.

import os, sys, re
import Tac, SecretCli
import errno

# Set the 'NOPDB' environment variable to prevent us from ever dropping into
# PDB, as doing so could permit a security violation.
os.environ[ 'NOPDB' ] = '1'
# Also set stdin to None for extra safety.
sys.stdin = None

startupConfig = '/mnt/flash/startup-config'
# If no startup-config exists, don't do anything
if not os.access( startupConfig, os.F_OK ):
   sys.exit( 0 )

# Find root password in the startup-config
encryptedPassword = '*'
sshKey = ''
rootSecretPlainTextRe = re.compile(
   r'\s*aaa +root +secret( +0)? +(.+)$' )
rootSecretMd5EncryptedRe = re.compile(
   r'\s*aaa +root +secret 5 (' + SecretCli.md5EncryptedPasswdRe + r')\s*$' )
rootSecretSha512EncryptedRe = re.compile(
   r'\s*aaa +root +secret sha512 (' + SecretCli.sha512EncryptedPasswdRe + r')\s*$' )
rootSecretScryptEncryptedRe = re.compile(
   r'\s*aaa +root +secret scrypt (' + SecretCli.scryptEncryptedPasswdRe + r')\s*$' )
rootSecretYescryptEncryptedRe = re.compile(
   r'\s*aaa +root +secret yescrypt (' +
   SecretCli.yescryptEncryptedPasswdRe + r')\s*$' )
rootNopasswordRe = re.compile( r'\s*aaa +root +nopassword\s*$' )
rootSshKeyRe = re.compile( r'\s*aaa +root +ssh-key( +0)? +(.+)$' )
rootNologinRe = re.compile( r'\s*no +aaa +root\s*$' )

# Use the sha512 algorithm to hash the password if entered in plaintext
SecretCli.defaultHashFunc = lambda: "sha512"

f = open( startupConfig ) # pylint: disable=consider-using-with
getPassword = False
for line in f:
   if not getPassword:
      m = rootSecretMd5EncryptedRe.match( line )
      if m:
         encryptedPassword = m.group( 1 )
         getPassword = True
         continue
      m = rootSecretSha512EncryptedRe.match( line )
      if m:
         encryptedPassword = m.group( 1 )
         getPassword = True
         continue
      m = rootSecretScryptEncryptedRe.match( line )
      if m:
         encryptedPassword = m.group( 1 )
         getPassword = True
         continue
      m = rootSecretYescryptEncryptedRe.match( line )
      if m:
         encryptedPassword = m.group( 1 )
         getPassword = True
         continue
      m = rootNopasswordRe.match( line )
      if m:
         encryptedPassword = ""
         getPassword = True
         continue
      m = rootNologinRe.match( line )
      if m:
         break
      # keep these last
      m = rootSshKeyRe.match( line )
      if m:
         # deal with spaces to be consistent with the CLI
         sshKey = re.sub( '  +', ' ', m.group( 2 ).strip() )
         # SshKey config should always be after root password
         # Finish the search if we find SshKey
         break
      m = rootSecretPlainTextRe.match( line )
      if m:
         # deal with spaces to be consistent with the CLI
         plaintext = re.sub( '  +', ' ', m.group( 2 ).strip() )
         encryptedPassword = SecretCli.setCleartextSecret( None, plaintext ).hash()
         getPassword = True
         continue
   else:
      m = rootSshKeyRe.match( line )
      if m:
         # deal with spaces to be consistent with the CLI
         sshKey = re.sub( '  +', ' ', m.group( 2 ).strip() )
      # Root ssh-key config should be right next to root password
      # Getting root password and the next line is not root ssh-key
      # means we only have root password
      break

passwordError = False
try:
   Tac.run( [ "/usr/sbin/usermod", "-p", encryptedPassword, "root" ] )
except: # pylint: disable-msg=W0702
   sys.stderr.write( "\nError in setting root password.\n" )
   # Nothing to do. root login is disabled( on shipping branches ) by default.
   # We need to set ssh key before exiting
   passwordError = True

if encryptedPassword != '*' and not passwordError:
   print( "Root login is enabled" )

filename = "/root/.ssh/authorized_keys"
sshKeyError = False
if sshKey:
   if not os.path.exists( os.path.dirname( filename ) ):
      try:
         os.makedirs( os.path.dirname( filename ) )
      except OSError as exc: # Guard against race condition
         if exc.errno != errno.EEXIST:
            sys.stderr.write( "\nError in creating root ssh key file.\n" )
            sshKeyError = True
   
   if not sshKeyError:
      with open( filename, "w+" ) as f:
         f.write( sshKey )
      print( "Root ssh key login is enabled" )

# Handle exception
if passwordError or sshKeyError:
   sys.exit( 1 )

sys.exit( 0 )
