<s> from __future__ import unicode_literals <EOL> import os <EOL> import shutil <EOL> import tempfile <EOL> import paramiko <EOL> from reviewboard . ssh . client import SSHClient <EOL> from reviewboard . ssh . errors import UnsupportedSSHKeyError <EOL> from reviewboard . ssh . storage import FileSSHStorage <EOL> from reviewboard . testing . testcase import TestCase <EOL> class SSHTestCase ( TestCase ) : <EOL> def setUp ( self ) : <EOL> super ( SSHTestCase , self ) . setUp ( ) <EOL> self . old_home = os . getenv ( '<STR_LIT>' ) <EOL> self . tempdir = None <EOL> os . environ [ '<STR_LIT>' ] = '<STR_LIT:0>' <EOL> FileSSHStorage . _ssh_dir = None <EOL> if not hasattr ( SSHTestCase , '<STR_LIT>' ) : <EOL> SSHTestCase . key1 = paramiko . RSAKey . generate ( <NUM_LIT> ) <EOL> SSHTestCase . key2 = paramiko . DSSKey . generate ( <NUM_LIT> ) <EOL> SSHTestCase . key1_b64 = SSHTestCase . key1 . get_base64 ( ) <EOL> SSHTestCase . key2_b64 = SSHTestCase . key2 . get_base64 ( ) <EOL> def tearDown ( self ) : <EOL> super ( SSHTestCase , self ) . tearDown ( ) <EOL> self . _set_home ( self . old_home ) <EOL> if self . tempdir : <EOL> shutil . rmtree ( self . tempdir ) <EOL> def _set_home ( self , homedir ) : <EOL> os . environ [ '<STR_LIT>' ] = homedir <EOL> class FileSSHStorageTests ( SSHTestCase ) : <EOL> """<STR_LIT>""" <EOL> def setUp ( self ) : <EOL> super ( FileSSHStorageTests , self ) . setUp ( ) <EOL> self . tempdir = tempfile . mkdtemp ( prefix = '<STR_LIT>' ) <EOL> self . _set_home ( self . tempdir ) <EOL> def test_get_ssh_dir_with_dot_ssh ( self ) : <EOL> """<STR_LIT>""" <EOL> sshdir = os . path . join ( self . tempdir , '<STR_LIT>' ) <EOL> storage = FileSSHStorage ( ) <EOL> self . assertEqual ( storage . get_ssh_dir ( ) , sshdir ) <EOL> def test_get_ssh_dir_with_ssh ( self ) : <EOL> """<STR_LIT>""" <EOL> sshdir = os . path . join ( self . tempdir , '<STR_LIT>' ) <EOL> os . mkdir ( sshdir , <NUM_LIT> ) <EOL> storage = FileSSHStorage ( ) <EOL> self . assertEqual ( storage . get_ssh_dir ( ) , sshdir ) <EOL> def test_get_ssh_dir_with_dot_ssh_and_localsite ( self ) : <EOL> """<STR_LIT>""" <EOL> sshdir = os . path . join ( self . tempdir , '<STR_LIT>' , '<STR_LIT>' ) <EOL> storage = FileSSHStorage ( namespace = '<STR_LIT>' ) <EOL> self . assertEqual ( storage . get_ssh_dir ( ) , sshdir ) <EOL> def test_get_ssh_dir_with_ssh_and_localsite ( self ) : <EOL> """<STR_LIT>""" <EOL> sshdir = os . path . join ( self . tempdir , '<STR_LIT>' ) <EOL> os . mkdir ( sshdir , <NUM_LIT> ) <EOL> sshdir = os . path . join ( sshdir , '<STR_LIT>' ) <EOL> storage = FileSSHStorage ( namespace = '<STR_LIT>' ) <EOL> self . assertEqual ( storage . get_ssh_dir ( ) , sshdir ) <EOL> def test_write_user_key_unsupported ( self ) : <EOL> """<STR_LIT>""" <EOL> class FakeKey ( object ) : <EOL> pass <EOL> storage = FileSSHStorage ( ) <EOL> self . assertRaises ( UnsupportedSSHKeyError , <EOL> lambda : storage . write_user_key ( FakeKey ( ) ) ) <EOL> def test_read_host_keys ( self ) : <EOL> """<STR_LIT>""" <EOL> storage = FileSSHStorage ( ) <EOL> storage . ensure_ssh_dir ( ) <EOL> line1 = '<STR_LIT>' % self . key1_b64 <EOL> line2 = '<STR_LIT>' % self . key2_b64 <EOL> filename = storage . get_host_keys_filename ( ) <EOL> with open ( filename , '<STR_LIT:w>' ) as fp : <EOL> fp . write ( '<STR_LIT>' % line1 ) <EOL> fp . write ( '<STR_LIT:\n>' ) <EOL> fp . write ( '<STR_LIT>' ) <EOL> fp . write ( '<STR_LIT>' % line2 ) <EOL> lines = storage . read_host_keys ( ) <EOL> self . assertEqual ( len ( lines ) , <NUM_LIT:2> ) <EOL> self . assertEqual ( lines [ <NUM_LIT:0> ] , line1 ) <EOL> self . assertEqual ( lines [ <NUM_LIT:1> ] , line2 ) <EOL> def test_add_host_key ( self ) : <EOL> """<STR_LIT>""" <EOL> storage = FileSSHStorage ( ) <EOL> storage . add_host_key ( '<STR_LIT>' , self . key1 ) <EOL> filename = storage . get_host_keys_filename ( ) <EOL> with open ( filename , '<STR_LIT:r>' ) as fp : <EOL> lines = fp . readlines ( ) <EOL> self . assertEqual ( len ( lines ) , <NUM_LIT:1> ) <EOL> self . assertEqual ( lines [ <NUM_LIT:0> ] , '<STR_LIT>' % self . key1_b64 ) <EOL> def test_replace_host_key ( self ) : <EOL> """<STR_LIT>""" <EOL> storage = FileSSHStorage ( ) <EOL> storage . add_host_key ( '<STR_LIT>' , self . key1 ) <EOL> storage . replace_host_key ( '<STR_LIT>' , self . key1 , self . key2 ) <EOL> filename = storage . get_host_keys_filename ( ) <EOL> with open ( filename , '<STR_LIT:r>' ) as fp : <EOL> lines = fp . readlines ( ) <EOL> self . assertEqual ( len ( lines ) , <NUM_LIT:1> ) <EOL> self . assertEqual ( lines [ <NUM_LIT:0> ] , '<STR_LIT>' % self . key2_b64 ) <EOL> def test_replace_host_key_no_known_hosts ( self ) : <EOL> """<STR_LIT>""" <EOL> storage = FileSSHStorage ( ) <EOL> storage . replace_host_key ( '<STR_LIT>' , self . key1 , self . key2 ) <EOL> filename = storage . get_host_keys_filename ( ) <EOL> with open ( filename , '<STR_LIT:r>' ) as fp : <EOL> lines = fp . readlines ( ) <EOL> self . assertEqual ( len ( lines ) , <NUM_LIT:1> ) <EOL> self . assertEqual ( lines [ <NUM_LIT:0> ] , '<STR_LIT>' % self . key2_b64 ) <EOL> class SSHClientTests ( SSHTestCase ) : <EOL> """<STR_LIT>""" <EOL> def setUp ( self ) : <EOL> super ( SSHClientTests , self ) . setUp ( ) <EOL> self . tempdir = tempfile . mkdtemp ( prefix = '<STR_LIT>' ) <EOL> def test_generate_user_key ( self , namespace = None ) : <EOL> """<STR_LIT>""" <EOL> self . _set_home ( self . tempdir ) <EOL> client = SSHClient ( namespace = namespace ) <EOL> key = client . generate_user_key ( bits = <NUM_LIT> ) <EOL> key_file = os . path . join ( client . storage . get_ssh_dir ( ) , '<STR_LIT>' ) <EOL> self . assertTrue ( os . path . exists ( key_file ) ) <EOL> self . assertEqual ( client . get_user_key ( ) , key ) <EOL> def test_generate_user_key_with_localsite ( self ) : <EOL> """<STR_LIT>""" <EOL> self . test_generate_user_key ( '<STR_LIT>' ) <EOL> def test_delete_user_key ( self , namespace = None ) : <EOL> """<STR_LIT>""" <EOL> self . _set_home ( self . tempdir ) <EOL> client = SSHClient ( namespace = namespace ) <EOL> client . import_user_key ( self . key1 ) <EOL> key_file = os . path . join ( client . storage . get_ssh_dir ( ) , '<STR_LIT>' ) <EOL> self . assertTrue ( os . path . exists ( key_file ) ) <EOL> self . assertEqual ( client . get_user_key ( ) , self . key1 ) <EOL> client . delete_user_key ( ) <EOL> self . assertFalse ( os . path . exists ( key_file ) ) <EOL> def test_delete_user_key_with_localsite ( self ) : <EOL> """<STR_LIT>""" <EOL> self . test_delete_user_key ( '<STR_LIT>' ) <EOL> def test_add_host_key ( self , namespace = None ) : <EOL> """<STR_LIT>""" <EOL> self . _set_home ( self . tempdir ) <EOL> client = SSHClient ( namespace = namespace ) <EOL> client . add_host_key ( '<STR_LIT>' , self . key1 ) <EOL> known_hosts_file = client . storage . get_host_keys_filename ( ) <EOL> self . assertTrue ( os . path . exists ( known_hosts_file ) ) <EOL> with open ( known_hosts_file , '<STR_LIT:r>' ) as f : <EOL> lines = f . readlines ( ) <EOL> self . assertEqual ( len ( lines ) , <NUM_LIT:1> ) <EOL> self . assertEqual ( lines [ <NUM_LIT:0> ] . split ( ) , <EOL> [ '<STR_LIT>' , self . key1 . get_name ( ) , self . key1_b64 ] ) <EOL> def test_add_host_key_with_localsite ( self ) : <EOL> """<STR_LIT>""" <EOL> self . test_add_host_key ( '<STR_LIT>' ) <EOL> def test_replace_host_key ( self , namespace = None ) : <EOL> """<STR_LIT>""" <EOL> self . _set_home ( self . tempdir ) <EOL> client = SSHClient ( namespace = namespace ) <EOL> client . add_host_key ( '<STR_LIT>' , self . key1 ) <EOL> client . replace_host_key ( '<STR_LIT>' , self . key1 , self . key2 ) <EOL> known_hosts_file = client . storage . get_host_keys_filename ( ) <EOL> self . assertTrue ( os . path . exists ( known_hosts_file ) ) <EOL> with open ( known_hosts_file , '<STR_LIT:r>' ) as f : <EOL> lines = f . readlines ( ) <EOL> self . assertEqual ( len ( lines ) , <NUM_LIT:1> ) <EOL> self . assertEqual ( lines [ <NUM_LIT:0> ] . split ( ) , <EOL> [ '<STR_LIT>' , self . key2 . get_name ( ) , <EOL> self . key2_b64 ] ) <EOL> def test_replace_host_key_with_localsite ( self ) : <EOL> """<STR_LIT>""" <EOL> self . test_replace_host_key ( '<STR_LIT>' ) <EOL> def test_import_user_key ( self , namespace = None ) : <EOL> """<STR_LIT>""" <EOL> self . _set_home ( self . tempdir ) <EOL> client = SSHClient ( namespace = namespace ) <EOL> client . import_user_key ( self . key1 ) <EOL> self . assertEqual ( client . get_user_key ( ) , self . key1 ) <EOL> def test_import_user_key_with_localsite ( self ) : <EOL> """<STR_LIT>""" <EOL> self . test_import_user_key ( '<STR_LIT>' ) </s>
