1import copy
2from unittest.mock import patch, MagicMock
3import pytest
4import os
5import random
6import tempfile
7
8from UM.Trust import TrustBasics, Trust
9
10from scripts.signfile import signFile
11from scripts.signfolder import signFolder
12
13_folder_names = ["a", "b"]
14_subfolder_names = ["sub", "."]
15_file_names = ["x.txt", "y.txt", "z.txt"]
16_passphrase = "swordfish"  # For code coverage: Securely storing a private key without one is probably better.
17
18
19class TestTrust:
20
21    # NOTE: Exhaustively testing trust is going to be difficult. We rely on audits (as well) in this matter.
22
23    @pytest.fixture()
24    def init_trust(self):
25        # Create a temporary directory and save a test key-pair to it:
26        temp_dir = tempfile.TemporaryDirectory()
27        temp_path = temp_dir.name
28        private_key, public_key = TrustBasics.generateNewKeyPair()
29        private_path = os.path.join(temp_path, "test_private_key.pem")
30        public_path = os.path.join(temp_path, "test_public_key.pem")
31        TrustBasics.saveKeyPair(private_key, private_path, public_path, _passphrase)
32
33        # Create random files:
34        all_paths = [os.path.abspath(os.path.join(temp_path, x, y, z))
35                     for x in _folder_names for y in _subfolder_names for z in _file_names]
36        for path in all_paths:
37            folder_path = os.path.dirname(path)
38            if not os.path.exists(folder_path):
39                os.makedirs(folder_path)
40            with open(path, "w") as file:
41                file.write("".join(random.choice(['a', 'b', 'c', '0', '1', '2', '\n']) for _ in range(1024)))
42
43        # Instantiate a trust object with the public key that was just generated:
44        violation_callback = MagicMock()
45        trust = Trust(public_path)  # No '.getInstance', since key & handler provided.
46        trust._violation_handler = violation_callback
47        yield temp_path, private_path, trust, violation_callback
48
49        temp_dir.cleanup()
50
51    def test_signFileAndVerify(self, init_trust):
52        temp_dir, private_path, trust_instance, violation_callback = init_trust
53        filepath_signed = os.path.join(temp_dir, _folder_names[0], _subfolder_names[0], _file_names[0])
54        filepath_unsigned = os.path.join(temp_dir, _folder_names[1], _subfolder_names[0], _file_names[2])
55
56        # Attempt to sign a file.
57        assert signFile(private_path, filepath_signed, _passphrase)
58
59        # Check if we're able to verify the file we just signed.
60        assert trust_instance.signedFileCheck(filepath_signed)
61        assert violation_callback.call_count == 0  # No violation
62
63        # Check if the file we didn't sign notifies us about this.
64        assert not trust_instance.signedFileCheck(filepath_unsigned)
65        assert violation_callback.call_count == 1
66
67        # An unknown file is also seen as an invalid one.
68        assert not trust_instance.signedFileCheck("file-not-found-check")
69        assert violation_callback.call_count == 2
70
71        # The signing should fail if we disable the key (since we can't confirm anything)
72        public_key = copy.copy(trust_instance._public_key)
73        trust_instance._public_key = None
74        assert not trust_instance.signedFileCheck(filepath_signed)
75        assert violation_callback.call_count == 3
76        violation_callback.reset_mock()
77        trust_instance._public_key = public_key
78
79        # Oh noes! Someone changed the file!
80        with open(filepath_signed, "w") as file:
81            file.write("\nPay 10 Golden Talents To Get Your Data Back Or Else\n")
82        assert not trust_instance.signedFolderCheck(filepath_signed)
83        assert violation_callback.call_count == 1
84        violation_callback.reset_mock()
85
86        # If one file is missing, the entire folder isn't considered to be signed.
87        os.remove(filepath_signed)
88        assert not trust_instance.signedFolderCheck(filepath_signed)
89        assert violation_callback.call_count == 1
90        violation_callback.reset_mock()
91
92    def test_signFolderAndVerify(self, init_trust):
93        temp_dir, private_path, trust_instance, violation_callback = init_trust
94        folderpath_signed = os.path.join(temp_dir, _folder_names[0])
95        folderpath_unsigned = os.path.join(temp_dir, _folder_names[1])
96
97        # Attempt to sign a folder & validate it's signatures.
98        assert signFolder(private_path, folderpath_signed, [], _passphrase)
99        assert trust_instance.signedFolderCheck(folderpath_signed)
100
101        # A folder that is not signed should be seen as such
102        assert not trust_instance.signedFolderCheck(folderpath_unsigned)
103        assert violation_callback.call_count == 1
104        violation_callback.reset_mock()
105
106        # Unknown folders should also be seen as unsigned
107        assert not trust_instance.signedFileCheck("folder-not-found-check")
108        assert violation_callback.call_count == 1
109        violation_callback.reset_mock()
110
111        # After removing the key, the folder that was signed should be seen as unsigned.
112        public_key = copy.copy(trust_instance._public_key)
113        trust_instance._public_key = None
114        assert not trust_instance.signedFolderCheck(folderpath_signed)
115        assert violation_callback.call_count == 1
116        violation_callback.reset_mock()
117        trust_instance._public_key = public_key
118
119        # Any modification will should also invalidate it.
120        filepath = os.path.join(folderpath_signed, _subfolder_names[0], _file_names[1])
121        with open(filepath, "w") as file:
122            file.write("\nAlice and Bob will never notice this! Hehehehe.\n")
123        assert not trust_instance.signedFolderCheck(folderpath_signed)
124        assert violation_callback.call_count > 0
125        violation_callback.reset_mock()
126
127        os.remove(filepath)
128        assert not trust_instance.signedFolderCheck(folderpath_signed)
129        assert violation_callback.call_count == 1
130        violation_callback.reset_mock()
131
132    def test_initTrustFail(self):
133        with pytest.raises(Exception):
134            Trust("key-not-found")
135
136        with pytest.raises(Exception):
137            Trust.getInstance()
138
139        assert Trust.getInstanceOrNone() is None
140
141    def test_keyIOFails(self):
142        private_key, public_key = TrustBasics.generateNewKeyPair()
143        assert not TrustBasics.saveKeyPair(private_key, public_key, "file-not-found", _passphrase)
144        assert TrustBasics.loadPrivateKey("key-not-found", _passphrase) is None
145
146    def test_signNonexisting(self):
147        private_key, public_key = TrustBasics.generateNewKeyPair()
148        assert TrustBasics.getFileSignature("file-not-found", private_key) is None
149