Skip to content

Commit 9f7796c

Browse files
leonav-unizarleonav-unizar
committed
Fix on inheritance of key comparison
1 parent f8dd379 commit 9f7796c

File tree

6 files changed

+30
-24
lines changed

6 files changed

+30
-24
lines changed

‎include/craper/key.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,10 +107,11 @@ class Key {
107107

108108
virtual error_handling::ProgramResult ExportKeyAsBinary(std::string out_file);
109109

110-
bool operator==(const std::shared_ptr<Key>& other) const;
110+
bool operator==(const Key& other) const;
111111

112112
struct KeyHashFunction {
113113
size_t operator()(const std::shared_ptr<Key>& key) const;
114+
bool operator()(const std::shared_ptr<Key>& lhs, const std::shared_ptr<Key>& rhs) const;
114115
};
115116

116117
protected:

‎include/craper/key_scanner.h‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -38,20 +38,20 @@ class ScannerFacade {
3838

3939
// Query
4040
bool IsProcessAlive() const;
41-
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction> DoScan();
42-
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction> GetKeys();
41+
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction> DoScan();
42+
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction> GetKeys();
4343
error_handling::ProgramResult ExportKeysToJSON(std::string output_json);
4444
error_handling::ProgramResult ExportKeysToBinary();
4545

4646
// Strategies
4747
void AddScanners(ScannerVector scanners);
4848

4949
private:
50-
void AddKeys(std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction> keys);
50+
void AddKeys(std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction> keys);
5151

5252
process_manipulation::ProcessCapturer capturer_;
5353
ScannerVector scanners_;
54-
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction> keys_;
54+
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction> keys_;
5555
OnDestroyAction on_destroy_;
5656
DWORD pid_;
5757
};

‎include/craper/scanners.h‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ namespace key_scanner {
1414
class ScanStrategy {
1515
public:
1616
virtual ~ScanStrategy() = default;
17-
virtual std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction> Scan(unsigned char* input_buffer, process_manipulation::HeapInformation heap_info) const = 0;
17+
virtual std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction> Scan(unsigned char* input_buffer, process_manipulation::HeapInformation heap_info) const = 0;
1818

1919
// iostream output
2020
virtual std::string GetName() const { return "Unnamed scanner"; };
@@ -23,7 +23,7 @@ class ScanStrategy {
2323
class CryptoAPIScan : public ScanStrategy {
2424
public:
2525
CryptoAPIScan() = default;
26-
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction> Scan(unsigned char* input_buffer, process_manipulation::HeapInformation heap_info) const override;
26+
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction> Scan(unsigned char* input_buffer, process_manipulation::HeapInformation heap_info) const override;
2727

2828
static void InitializeCryptoAPI();
2929
std::string GetName() const override { return "CryptoAPI Key Scanner"; };
@@ -37,7 +37,7 @@ class CryptoAPIScan : public ScanStrategy {
3737
class RoundKeyScan : public ScanStrategy {
3838
public:
3939
RoundKeyScan() = default;
40-
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction> Scan(unsigned char* input_buffer, process_manipulation::HeapInformation heap_info) const override;
40+
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction> Scan(unsigned char* input_buffer, process_manipulation::HeapInformation heap_info) const override;
4141

4242
std::string GetName() const override { return "AES Round Key Scanner"; };
4343
};

‎src/key.cc‎

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -66,12 +66,11 @@ error_handling::ProgramResult Key::ExportKeyAsBinary(std::string out_file) {
6666
return OkResult("Successfully exported key blob to " + out_file);
6767
}
6868

69-
bool Key::operator==(const std::shared_ptr<Key>& other) const {
69+
bool Key::operator==(const Key& other) const {
70+
if (this->GetAlgorithm() != other.GetAlgorithm()) return false;
71+
if (this->GetSize() != other.GetSize()) return false;
7072

71-
if (this->GetSize() != other->GetSize())
72-
return false;
73-
74-
return (this->GetKey() == other->GetKey());
73+
return (this->GetKey() == other.GetKey());
7574
}
7675

7776
std::size_t KeyType::GetSize() const {
@@ -93,7 +92,7 @@ std::string KeyType::GetAlgorithmAsString() const {
9392

9493

9594
size_t Key::KeyHashFunction::operator()(const std::shared_ptr<Key>& key) const {
96-
const vector<unsigned char> key_data = key->GetKey();
95+
const std::vector<unsigned char> key_data = key->GetKey();
9796
if (key_data.empty()) {
9897
return 0;
9998
}
@@ -102,9 +101,15 @@ size_t Key::KeyHashFunction::operator()(const std::shared_ptr<Key>& key) const {
102101
for (size_t i = 1; i < key_data.size(); ++i) {
103102
hash = hash ^ std::hash<unsigned char>()(key_data[i]);
104103
}
104+
105105
return hash;
106106
}
107107

108+
bool Key::KeyHashFunction::operator()(const std::shared_ptr<Key>& lhs, const std::shared_ptr<Key>& rhs) const {
109+
if (!lhs || !rhs) return !lhs && !rhs; // Handle null pointers
110+
return *lhs == *rhs;
111+
}
112+
108113

109114
vector<BYTE> CrAPIKeyWrapper::GetParameter(DWORD parameter) {
110115
vector<BYTE> data_bytes;
@@ -183,11 +188,11 @@ bool CryptoAPIKey::IsAsymmetricAlgorithm() {
183188

184189
error_handling::ProgramResult CryptoAPIKey::ExportKeyAsBinary(std::string out_file) {
185190
if (IsSymmetricAlgorithm()) {
186-
printf(" [!] Symmetric algorithm detected. Exporting the key as PLAINTEXTKEYBLOB");
191+
printf(" [!] Symmetric algorithm detected. Exporting the key as PLAINTEXTKEYBLOB\n");
187192
return ExportAsBinaryGeneric(PLAINTEXTKEYBLOB, out_file + ".bin");
188193

189194
} else if (IsAsymmetricAlgorithm()) {
190-
printf(" [!] An asymmetric key was detected. Since it's unkown whether it is the private or public pair, it will be exported in both formats");
195+
printf(" [!] An asymmetric key was detected. Since it's unkown whether it is the private or public pair, it will be exported in both formats\n");
191196
ProgramResult pr = ExportAsBinaryGeneric(PUBLICKEYBLOB, out_file + ".PUBK");
192197
if (pr.IsErr()) return pr;
193198
pr = ExportAsBinaryGeneric(PRIVATEKEYBLOB, out_file + ".PRIVK");

‎src/key_scanner.cc‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ bool ScannerFacade::IsProcessAlive() const {
3939
return capturer_.IsProcessAlive();
4040
}
4141

42-
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction> ScannerFacade::DoScan() {
42+
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction> ScannerFacade::DoScan() {
4343
printf(" [i] Starting scan\n");
4444
if (scanners_.size() == 0) {
4545
printf(" [x] No scanner selected\n");
@@ -81,7 +81,7 @@ std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction> ScannerFacade::Do
8181
return keys_;
8282
}
8383

84-
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction> ScannerFacade::GetKeys() {
84+
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction> ScannerFacade::GetKeys() {
8585
return keys_;
8686
}
8787

@@ -126,7 +126,7 @@ void ScannerFacade::AddScanners(ScannerVector scanners) {
126126
scanners.clear();
127127
}
128128

129-
void ScannerFacade::AddKeys(std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction> keys) {
129+
void ScannerFacade::AddKeys(std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction> keys) {
130130
keys_.merge(keys);
131131
}
132132

‎src/scanners.cc‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -78,9 +78,9 @@ void CryptoAPIScan::InitializeCryptoAPI() {
7878
}
7979
}
8080
}
81-
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction> CryptoAPIScan::Scan(unsigned char *input_buffer, HeapInformation heap_info) const {
81+
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction> CryptoAPIScan::Scan(unsigned char *input_buffer, HeapInformation heap_info) const {
8282

83-
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction> found_keys = std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction>();
83+
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction> found_keys = std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction>();
8484

8585
InitializeCryptoAPI();
8686
if (cryptoapi_functions_initialized) {
@@ -137,7 +137,7 @@ std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction> CryptoAPIScan::Sc
137137

138138
if (heap_info.RebaseAddress(&ptr, (ULONG_PTR) input_buffer)) {
139139
found_keys.insert(
140-
std::make_unique<CryptoAPIKey>(
140+
std::make_shared<CryptoAPIKey>(
141141
CryptoAPIKey(key_data_struct, (unsigned char*) ptr)
142142
)
143143
);
@@ -171,8 +171,8 @@ std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction> CryptoAPIScan::Sc
171171
return found_keys;
172172
}
173173

174-
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction> RoundKeyScan::Scan(unsigned char *buffer, HeapInformation heap_info) const {
175-
return std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction>();
174+
std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction> RoundKeyScan::Scan(unsigned char *buffer, HeapInformation heap_info) const {
175+
return std::unordered_set<std::shared_ptr<Key>, Key::KeyHashFunction, Key::KeyHashFunction>();
176176
}
177177

178178
ScannerVector::ScannerVector(std::unique_ptr<std::vector<std::unique_ptr<ScanStrategy>>> scanners) {

0 commit comments

Comments
 (0)