address revioew feedback Azure#2

yugangw-msft · Oct 22, 2016 · 70138bb · 70138bb
1 parent 295f9b2
commit 70138bb
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_actions.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_actions.py
@@ -278,6 +278,7 @@ def _generate_ssh_keys(private_key_filepath, public_key_filepath):
     ssh_dir, _ = os.path.split(private_key_filepath)
     if not os.path.exists(ssh_dir):
         os.makedirs(ssh_dir)
+        os.chmod(ssh_dir, 0o700)
 
     key = paramiko.RSAKey.generate(2048)
     key.write_private_key_file(private_key_filepath)
@@ -286,6 +287,7 @@ def _generate_ssh_keys(private_key_filepath, public_key_filepath):
     with open(public_key_filepath, 'w') as public_key_file:
         public_key = '%s %s'  % (key.get_name(), key.get_base64())
         public_key_file.write(public_key)
+    os.chmod(public_key_filepath, 0o644)
 
     return public_key
 

diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/test_vm_actions.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/test_vm_actions.py
@@ -46,3 +46,13 @@ def test_generate_specfied_ssh_key_files(self):
         #still a file name
         self.assertEqual(args3.ssh_key_value, public_key_file2)
 
+        #4 verify file naming if the pub file doesn't end with .pub
+        _, public_key_file4 = tempfile.mkstemp()
+        public_key_file4 += '1' #make it nonexisting
+        args4 = mock.MagicMock()
+        args4.ssh_key_value = public_key_file4
+        args4.generate_ssh_keys = True
+        _handle_container_ssh_file(command='acs create', args=args4)
+        self.assertTrue(os.path.isfile(public_key_file4 + '.private'))
+        self.assertTrue(os.path.isfile(public_key_file4))
+