add hosts entries and ssh keys only once

Some users add a newline to their key, strip the key when checking if
it's already added.

Stop using regex (that didn't work) when checking if we already added a
hosts entry.  Use a simple startswith() instead.

Switch to use MagicMock as it has support for iteration.

Change-Id: Ibeabbf1c4448933cdf274e6667c5a9493f3e5e37
Signed-off-by: Matthew Thode <mthode@mthode.org>
This commit is contained in:
Matthew Thode 2020-12-30 15:31:24 -06:00
parent a507b71117
commit 4ef3e3b0e7
No known key found for this signature in database
GPG Key ID: 64A37BEAAE19A4E8
2 changed files with 9 additions and 10 deletions

View File

@ -22,7 +22,6 @@ import errno
import json
import logging
import os
import re
import subprocess
import sys
import time
@ -1419,7 +1418,7 @@ def write_ssh_keys(args):
if key_title not in keys_to_write:
keys_to_write.append(key_title)
if key not in keys_to_write:
if key.rstrip() not in keys_to_write:
keys_to_write.append(key)
files_to_write = {
@ -1492,14 +1491,14 @@ def set_hostname_from_config_drive(args):
for host in hosts_to_add:
host_value = hosts_to_add[host]
# See if we already have a hosts entry for hostname
prog = re.compile('^%s .*%s\n' % (host_value, host))
match = None
if os.path.isfile('/etc/hosts'):
with open('/etc/hosts') as fh:
match = prog.match(fh.read())
# Write out a hosts entry for hostname
if match is None:
with safe_open('/etc/hosts', 'r+') as fh:
for line in fh:
if line.startswith('%s %s' % (host_value, host)):
break
else:
fh.write(u'%s %s\n' % (host_value, host))
else:
with safe_open('/etc/hosts', 'a+') as fh:
fh.write(u'%s %s\n' % (host_value, host))

View File

@ -92,7 +92,7 @@ class TestGlean(base.BaseTestCase):
# note; don't use spec=file here ... it's not py3
# compatible. It really just limits the allowed
# mocked functions.
mock_handle = mock.Mock()
mock_handle = mock.MagicMock()
mock_handle.__enter__ = mock.Mock()
mock_handle.__exit__ = mock.Mock()
mock_handle.name = path