webssh/tests/test_storage.py

import os
import shutil
import tempfile
import unittest

from webssh.crypto import is_encrypted
from webssh.storage import ConnectionStore


SECRET = 'test-server-secret'


class TestConnectionStore(unittest.TestCase):

    def setUp(self):
        self.tmpdir = tempfile.mkdtemp()
        self.filename = os.path.join(self.tmpdir, 'connections.json')
        self.store = ConnectionStore(self.filename, SECRET)

    def tearDown(self):
        shutil.rmtree(self.tmpdir)

    def read_raw(self):
        with open(self.filename, encoding='utf-8') as f:
            return f.read()

    def test_upsert_list_and_delete(self):
        profile = self.store.upsert('admin', {
            'hostname': '127.0.0.1',
            'port': 22,
            'username': 'root',
            'password': 'root-secret',
            'term': 'xterm-256color',
            'directory': '/var/www',
            'auth_type': 'password'
        })

        self.assertEqual(profile['title'], 'root@127.0.0.1:22')
        self.assertEqual(profile['directory'], '/var/www')
        self.assertEqual([profile], self.store.list('admin'))
        self.assertEqual([], self.store.list('other'))

        updated = self.store.upsert('admin', {
            'hostname': '127.0.0.1',
            'port': 22,
            'username': 'root',
            'password': 'new-secret',
            'term': 'xterm',
            'auth_type': 'privatekey'
        })

        self.assertEqual(profile['id'], updated['id'])
        self.assertEqual(1, len(self.store.list('admin')))
        self.assertEqual('xterm', self.store.list('admin')[0]['term'])
        self.assertEqual('', self.store.list('admin')[0]['directory'])
        self.assertTrue(self.store.delete('admin', profile['id']))
        self.assertFalse(self.store.delete('admin', profile['id']))
        self.assertEqual([], self.store.list('admin'))

    def test_password_never_leaves_the_store(self):
        profile = self.store.upsert('admin', {
            'hostname': '127.0.0.1',
            'username': 'root',
            'password': 'root-secret'
        })

        self.assertNotIn('password', profile)
        self.assertNotIn('password_enc', profile)
        self.assertTrue(profile['has_password'])

        listed = self.store.list('admin')[0]
        self.assertNotIn('password', listed)
        self.assertNotIn('password_enc', listed)
        self.assertTrue(listed['has_password'])

        # ... but it is still usable server side
        self.assertEqual(
            'root-secret', self.store.get_password('admin', profile['id'])
        )

    def test_password_is_encrypted_at_rest(self):
        self.store.upsert('admin', {
            'hostname': '127.0.0.1',
            'username': 'root',
            'password': 'root-secret'
        })

        raw = self.read_raw()
        self.assertNotIn('root-secret', raw)
        self.assertIn('password_enc', raw)

    def test_password_is_scoped_to_owner_and_destination(self):
        profile = self.store.upsert('admin', {
            'hostname': '127.0.0.1',
            'username': 'root',
            'password': 'root-secret'
        })

        self.assertEqual(
            '', self.store.get_password('mallory', profile['id'])
        )
        self.assertEqual('', self.store.get_password('admin', 'deadbeef'))

        # an edited hostname yields a different id, so nothing is found
        other_id = self.store.make_id('admin', 'evil.example.com', 22, 'root')
        self.assertNotEqual(profile['id'], other_id)
        self.assertEqual('', self.store.get_password('admin', other_id))

    def test_profile_without_password(self):
        profile = self.store.upsert('admin', {
            'hostname': '127.0.0.1',
            'username': 'root'
        })
        self.assertFalse(profile['has_password'])
        self.assertEqual('', self.store.get_password('admin', profile['id']))

    def test_a_wrong_secret_does_not_break_the_store(self):
        profile = self.store.upsert('admin', {
            'hostname': '127.0.0.1',
            'username': 'root',
            'password': 'root-secret'
        })

        rotated = ConnectionStore(self.filename, 'a-different-secret')
        self.assertEqual('', rotated.get_password('admin', profile['id']))
        self.assertTrue(rotated.list('admin')[0]['has_password'])

    def test_encrypt_plaintext_passwords_migrates_old_files(self):
        legacy = ConnectionStore(self.filename)
        profile = legacy.upsert('admin', {
            'hostname': '127.0.0.1',
            'username': 'root',
            'password': 'legacy-secret'
        })
        self.assertIn('legacy-secret', self.read_raw())
        self.assertEqual(
            'legacy-secret', legacy.get_password('admin', profile['id'])
        )

        self.assertEqual(1, self.store.encrypt_plaintext_passwords())

        raw = self.read_raw()
        self.assertNotIn('legacy-secret', raw)
        self.assertEqual(
            'legacy-secret', self.store.get_password('admin', profile['id'])
        )
        self.assertTrue(self.store.list('admin')[0]['has_password'])

        # already migrated, nothing left to do
        self.assertEqual(0, self.store.encrypt_plaintext_passwords())

    def test_directory_is_normalized_and_optional(self):
        profile = self.store.upsert('admin', {
            'hostname': 'example.com',
            'username': 'deploy',
            'directory': '  /srv/app  '
        })
        self.assertEqual('/srv/app', profile['directory'])

        profile = self.store.upsert('admin', {
            'hostname': 'example.org',
            'username': 'deploy'
        })
        self.assertEqual('', profile['directory'])

    def test_privatekey_is_encrypted_and_never_public(self):
        profile = self.store.upsert('admin', {
            'hostname': 'example.com', 'username': 'root',
            'privatekey': 'private-key-material',
            'passphrase': 'key-passphrase', 'save_passphrase': True,
            'auth_type': 'privatekey'
        })
        raw = self.read_raw()
        self.assertNotIn('private-key-material', raw)
        self.assertNotIn('key-passphrase', raw)
        self.assertTrue(profile['has_privatekey'])
        self.assertTrue(profile['has_passphrase'])
        for public in (profile, self.store.list('admin')[0]):
            for field in ('privatekey', 'privatekey_enc', 'passphrase',
                          'passphrase_enc'):
                self.assertNotIn(field, public)
        self.assertEqual(
            ('private-key-material', 'key-passphrase'),
            self.store.get_privatekey('admin', profile['id'])
        )
        self.assertEqual(
            ('', ''), self.store.get_privatekey('other', profile['id'])
        )
        other_id = self.store.make_id('admin', 'evil.example', 22, 'root')
        self.assertEqual(('', ''),
                         self.store.get_privatekey('admin', other_id))
        rotated = ConnectionStore(self.filename, 'other-secret')
        self.assertEqual(('', ''),
                         rotated.get_privatekey('admin', profile['id']))

    def test_credentials_can_be_replaced_and_cleared(self):
        connection = {
            'hostname': 'example.com', 'username': 'root',
            'password': 'password-secret', 'privatekey': 'old-key',
            'passphrase': 'key-secret', 'save_passphrase': True
        }
        profile = self.store.upsert('admin', connection)
        connection.update(privatekey='new-key', save_passphrase=False)
        self.store.upsert('admin', connection)
        self.assertEqual(('new-key', ''),
                         self.store.get_privatekey('admin', profile['id']))
        connection['save_credentials'] = False
        public = self.store.upsert('admin', connection)
        self.assertFalse(public['has_password'])
        self.assertFalse(public['has_privatekey'])
        self.assertFalse(public['has_passphrase'])
        self.assertEqual('', self.store.get_password('admin', profile['id']))
        self.assertEqual(('', ''),
                         self.store.get_privatekey('admin', profile['id']))

    def test_privatekey_requires_encryption_secret(self):
        store = ConnectionStore(self.filename)
        with self.assertRaises(ValueError):
            store.upsert('admin', {
                'hostname': 'example.com', 'username': 'root',
                'privatekey': 'private-key-material'
            })
        self.assertFalse(os.path.exists(self.filename))

    def test_metadata_update_keeps_credentials_and_owner_boundary(self):
        profile = self.store.upsert('admin', {
            'hostname': 'example.com', 'username': 'root',
            'password': 'secret', 'name': 'Production', 'group': 'Servers'
        })
        self.assertIsNone(self.store.update_metadata(
            'other', profile['id'], {'name': 'stolen'}
        ))
        updated = self.store.update_metadata('admin', profile['id'], {
            'name': 'New name', 'group': 'Testing', 'notes': 'A note'
        })
        self.assertEqual('New name', updated['name'])
        self.assertEqual('secret', self.store.get_password('admin', profile['id']))

    def test_plain_export_omits_credentials_and_import_preserves_them(self):
        profile = self.store.upsert('admin', {
            'hostname': 'example.com', 'username': 'root', 'password': 'secret'
        })
        exported = self.store.export_connections('admin')
        self.assertNotIn('password', exported[0])
        self.assertNotIn('owner', exported[0])
        exported[0].update(name='Imported', password='injected')
        self.store.import_connections('admin', exported)
        self.assertEqual('secret', self.store.get_password('admin', profile['id']))
        self.assertEqual('Imported', self.store.list('admin')[0]['name'])

    def test_invalid_import_is_atomic(self):
        with self.assertRaises(ValueError):
            self.store.import_connections('admin', [
                {'hostname': 'example.com', 'username': 'root', 'port': 22},
                {'hostname': 'example.com', 'username': 'root', 'port': -1}
            ])
        self.assertEqual([], self.store.list('admin'))
        with self.assertRaises(ValueError):
            self.store.import_connections('admin', [
                {'hostname': 'example.com', 'username': 'root', 'port': 0}
            ])

    def test_encrypted_import_reencrypts_for_destination_server(self):
        profile = self.store.upsert('admin', {
            'hostname': 'example.com', 'username': 'root',
            'privatekey': 'key-material', 'passphrase': 'key-secret',
            'save_passphrase': True
        })
        exported = self.store.export_connections('admin', True)
        destination = ConnectionStore(self.filename + '.new', 'new-server')
        destination.import_connections('new-owner', exported, True)
        new_id = destination.make_id('new-owner', 'example.com', 22, 'root')
        self.assertEqual(('key-material', 'key-secret'),
                         destination.get_privatekey('new-owner', new_id))
        with open(destination.filename, encoding='utf-8') as saved:
            self.assertNotIn('key-material', saved.read())
        self.assertEqual([], destination.list('admin'))


class TestSecretBox(unittest.TestCase):

    def test_round_trip(self):
        from webssh.crypto import SecretBox

        box = SecretBox(SECRET)
        token = box.encrypt('hunter2')
        self.assertTrue(is_encrypted(token))
        self.assertNotIn('hunter2', token)
        self.assertEqual('hunter2', box.decrypt(token))

    def test_nonce_is_not_reused(self):
        from webssh.crypto import SecretBox

        box = SecretBox(SECRET)
        self.assertNotEqual(box.encrypt('same'), box.encrypt('same'))

    def test_empty_and_tampered_values(self):
        from webssh.crypto import SecretBox

        box = SecretBox(SECRET)
        self.assertEqual('', box.encrypt(''))
        self.assertEqual('', box.decrypt(''))
        self.assertEqual('', box.decrypt('not-a-token'))
        self.assertEqual('', box.decrypt('aesgcm$####'))

        token = box.encrypt('hunter2')
        self.assertEqual('', box.decrypt(token[:-4] + 'AAAA'))
        self.assertEqual('', SecretBox('other').decrypt(token))