From 9c7fbd1a90ba91357f5378699328c7ac1d06dc46 Mon Sep 17 00:00:00 2001
From: Miroslav Stampar <miroslav.stampar@gmail.com>
Date: Sun, 6 Apr 2014 18:19:54 +0200
Subject: [PATCH] Minor refactoring

---
 lib/core/enums.py        |  5 +++++
 lib/takeover/registry.py | 13 +++++++------
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/lib/core/enums.py b/lib/core/enums.py
index 53fce1aff..bc87b5a2a 100644
--- a/lib/core/enums.py
+++ b/lib/core/enums.py
@@ -140,6 +140,11 @@ class PROXY_TYPE:
     SOCKS4 = "SOCKS4"
     SOCKS5 = "SOCKS5"
 
+class REGISTRY_OPERATION:
+    READ = "read"
+    ADD = "add"
+    DELETE = "delete"
+
 class DUMP_FORMAT:
     CSV = "CSV"
     HTML = "HTML"
diff --git a/lib/takeover/registry.py b/lib/takeover/registry.py
index ab66758d5..36dbe295f 100644
--- a/lib/takeover/registry.py
+++ b/lib/takeover/registry.py
@@ -10,6 +10,7 @@ import os
 from lib.core.common import randomStr
 from lib.core.data import conf
 from lib.core.data import logger
+from lib.core.enums import REGISTRY_OPERATION
 
 class Registry:
     """
@@ -49,11 +50,11 @@ class Registry:
     def _createLocalBatchFile(self):
         self._batPathFp = open(self._batPathLocal, "w")
 
-        if self._operation == "read":
+        if self._operation == REGISTRY_OPERATION.READ:
             lines = self._batRead
-        elif self._operation == "add":
+        elif self._operation == REGISTRY_OPERATION.ADD:
             lines = self._batAdd
-        elif self._operation == "delete":
+        elif self._operation == REGISTRY_OPERATION.DELETE:
             lines = self._batDel
 
         for line in lines:
@@ -70,7 +71,7 @@ class Registry:
         os.unlink(self._batPathLocal)
 
     def readRegKey(self, regKey, regValue, parse=False):
-        self._operation = "read"
+        self._operation = REGISTRY_OPERATION.READ
 
         Registry._initVars(self, regKey, regValue, parse=parse)
         self._createRemoteBatchFile()
@@ -90,7 +91,7 @@ class Registry:
         return data
 
     def addRegKey(self, regKey, regValue, regType, regData):
-        self._operation = "add"
+        self._operation = REGISTRY_OPERATION.ADD
 
         Registry._initVars(self, regKey, regValue, regType, regData)
         self._createRemoteBatchFile()
@@ -103,7 +104,7 @@ class Registry:
         self.delRemoteFile(self._batPathRemote)
 
     def delRegKey(self, regKey, regValue):
-        self._operation = "delete"
+        self._operation = REGISTRY_OPERATION.DELETE
 
         Registry._initVars(self, regKey, regValue)
         self._createRemoteBatchFile()