From 3eaf22b1a66860476e70b26097207045f20435ea Mon Sep 17 00:00:00 2001
From: Matt Soucy <msoucy@csh.rit.edu>
Date: Tue, 27 May 2014 02:38:48 -0400
Subject: [PATCH] Make use of context managers to safely close files

---
 gamespy/gs_utility.py | 13 ++++++-------
 nas_server.py         | 24 +++++++++++++++---------
 storage_server.py     | 10 ++++------
 3 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/gamespy/gs_utility.py b/gamespy/gs_utility.py
index 11f3f08..2134245 100644
--- a/gamespy/gs_utility.py
+++ b/gamespy/gs_utility.py
@@ -5,15 +5,14 @@ import time
 import other.utils as utils
 
 def generate_secret_keys(filename="gslist.cfg"):
-    key_file = open(filename)
-
     secret_key_list = {}
-    for line in key_file.readlines():
-        #name = line[:54].strip() # Probably won't do anything with the name for now.
-        id = line[54:54+19].strip()
-        key = line[54+19:].strip()
+    with open(filename) as key_file:
+        for line in key_file.readlines():
+            #name = line[:54].strip() # Probably won't do anything with the name for now.
+            id = line[54:54+19].strip()
+            key = line[54+19:].strip()
 
-        secret_key_list[id] = key
+            secret_key_list[id] = key
 
     return secret_key_list
 
diff --git a/nas_server.py b/nas_server.py
index e34d346..52ce9dc 100644
--- a/nas_server.py
+++ b/nas_server.py
@@ -154,6 +154,15 @@ class NasHTTPServerHandler(BaseHTTPServer.BaseHTTPRequestHandler):
             dlcpath = "dlc/" + post["gamecd"]
             dlc_contenttype = False
 
+            def safeloadfi(fn, mode='rb'):
+                '''
+                safeloadfi : string -> string
+
+                Safely load contents of a file, given a filename, and closing the file afterward
+                '''
+                with open(os.path.join(dlcpath, fn), mode) as fi:
+                    return fi.read()
+
             if action == "count":
                 if post["gamecd"] in gamecodes_return_random_file:
                     ret = "1"
@@ -173,10 +182,10 @@ class NasHTTPServerHandler(BaseHTTPServer.BaseHTTPRequestHandler):
                             attr3 = None
                             if "attr3" in post:
                                 attr3 = post["attr3"]
-                            list = open(dlcpath + "/_list.txt", "rb").read()
-                            list = self.filter_list(list, attr1, attr2, attr3)
 
-                            count = self.get_file_count(list)
+                            lst = safeloadfi("_list.txt")
+                            lst = self.filter_list(dlcfi.read(), attr1, attr2, attr3)
+                            count = self.get_file_count(lst)
 
                     ret = "%d" % count
 
@@ -197,9 +206,8 @@ class NasHTTPServerHandler(BaseHTTPServer.BaseHTTPRequestHandler):
                 if os.path.exists(dlcpath):
                     # Look for a list file first.
                     # If the list file exists, send the entire thing back to the client.
-                    if os.path.isfile(dlcpath + "/_list.txt"):
-                        ret = open(dlcpath + "/_list.txt", "rb").read()
-                        ret = self.filter_list(ret, attr1, attr2, attr3)
+                    if os.path.isfile(os.path.join(dlcpath, "_list.txt")):
+                        ret = self.filter_list(safeloadfi("_list.txt"), attr1, attr2, attr3)
 
                         if post["gamecd"] in gamecodes_return_random_file:
                             ret = self.filter_list_random_files(ret, 1)
@@ -208,9 +216,7 @@ class NasHTTPServerHandler(BaseHTTPServer.BaseHTTPRequestHandler):
                 # Get only the base filename just in case there is a path involved somewhere in the filename string.
                 dlc_contenttype = True
                 contents = os.path.basename(post["contents"])
-
-                if os.path.isfile(dlcpath + "/" + contents):
-                    ret = open(dlcpath + "/" + contents, "rb").read()
+                ret = safeloadfi(contents)
 
             self.send_response(200)
 
diff --git a/storage_server.py b/storage_server.py
index 11ea860..b2b1566 100644
--- a/storage_server.py
+++ b/storage_server.py
@@ -374,9 +374,8 @@ class StorageHTTPServerHandler(BaseHTTPServer.BaseHTTPRequestHandler):
                 path = userdir + '/' + str(fileid)
                 cursor.execute('UPDATE filepaths SET path = ? WHERE fileid = ?', (path, fileid))
                 
-                file = open(path, 'wb')
-                file.write(filedata['data'][0])
-                file.close()
+                with open(path, 'wb') as fi:
+                    fi.write(filedata['data'][0])
             else:
                 logger.log(logging.WARNING, "Tried to upload big file, rejected. (%s bytes)", len(filedata['data'][0]))
                 fileid = 0
@@ -419,9 +418,8 @@ class StorageHTTPServerHandler(BaseHTTPServer.BaseHTTPRequestHandler):
                     filename = cursor.fetchone()[0]
 
                     if os.path.exists(filename):
-                        file = open(filename, 'rb')
-                        ret = file.read()
-                        file.close()
+                        with open(filename, 'rb') as fi:
+                            ret = fi.read()
                     else:
                         logger.log(logging.ERROR, "User is trying to access file that should exist according to DB, but doesn't! (%s)", filename)
                 except: