Merge pull request #128 from hxer/add/update

add/pocsuite auto update from github

Author: Hysia

pocsuite/lib/core/common.py

@@ -16,6 +16,7 @@
 import posixpath
 import marshal
 import unicodedata
+import time
 from pocsuite.lib.core.data import conf
 from pocsuite.lib.core.convert import stdoutencode
 from pocsuite.lib.core.log import LOGGER_HANDLER
@@ -428,3 +429,27 @@ def reIndent(s, numSpace):
     leadingSpace = numSpace * ' '
     lines = [leadingSpace + line for line in s.splitlines()]
     return '\n'.join(lines)
+
+
+def poll_process(process, suppress_errors=False):
+    """
+    Checks for process status (prints . if still running)
+    """
+
+    while True:
+        dataToStdout(".")
+        time.sleep(1)
+
+        returncode = process.poll()
+
+        if returncode is not None:
+            if not suppress_errors:
+                if returncode == 0:
+                    dataToStdout(" done\n")
+                elif returncode < 0:
+                    dataToStdout(" process terminated by signal %d\n"
+                                 % returncode)
+                elif returncode > 0:
+                    dataToStdout(" quit unexpectedly with return code %d\n"
+                                 % returncode)
+            break

pocsuite/lib/core/update.py

@@ -0,0 +1,94 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+"""
+Copyright (c) 2014-2016 pocsuite developers (https://seebug.org)
+See the file 'docs/COPYING' for copying permission
+"""
+
+import locale
+import os
+import re
+import subprocess
+import time
+
+from pocsuite.lib.core.common import dataToStdout
+from pocsuite.lib.core.common import poll_process
+from pocsuite.lib.core.data import conf
+from pocsuite.lib.core.data import logger
+from pocsuite.lib.core.data import paths
+from pocsuite.lib.core.revision import getRevisionNumber
+from pocsuite.lib.core.settings import GIT_REPOSITORY
+from pocsuite.lib.core.settings import IS_WIN
+from pocsuite.lib.core.settings import SITE
+
+
+def update():
+    if not conf.update:
+        return
+
+    success = False
+
+    if not os.path.exists(os.path.join(
+                    os.path.dirname(paths.POCSUITE_ROOT_PATH), ".git")):
+        err_msg = ("not a git repository. Please checkout the 'pocsuite' "
+                   "repository from GitHub (e.g. 'git clone --depth 1"
+                   "https://github.com/knownsec/Pocsuite.git pocsuite')")
+        logger.error(err_msg)
+    else:
+        info_msg = ("updating pocsuite to the latest development version from "
+                    "the GitHub repository")
+        logger.info(info_msg)
+
+        debug_msg = "pocsuite will try to update itself using 'git' command"
+        logger.debug(debug_msg)
+
+        dataToStdout("\r[%s] [INFO] update in progress " % time.strftime("%X"))
+
+        try:
+            # Reference:
+            # http://blog.stastnarodina.com/honza-en/spot/python-unicodeencodeerror/
+            process = subprocess.Popen(
+                        "git checkout . && git pull %s HEAD" % GIT_REPOSITORY,
+                        shell=True,
+                        stdout=subprocess.PIPE,
+                        stderr=subprocess.PIPE,
+                        cwd=paths.POCSUITE_ROOT_PATH.encode(
+                                locale.getpreferredencoding())
+                    )
+            poll_process(process, True)
+            stdout, stderr = process.communicate()
+            success = not process.returncode
+        except (IOError, OSError) as ex:
+            success = False
+            stderr = "{}".format(ex)
+
+        if success:
+            info_msg = "{0} the latest revision '{1}'".format(
+                        "already at" if "Already" in stdout else "updated to",
+                        getRevisionNumber())
+            logger.info(info_msg)
+        else:
+            if "Not a git repository" in stderr:
+                err_msg = ("not a valid git repository. Please checkout the "
+                           "'pocsuite' repository from GitHub "
+                           "(e.g. 'git clone --depth 1 "
+                           "https://github.com/knownsec/Pocsuite.git "
+                           "pocsuite')")
+                logger.error(err_msg)
+            else:
+                err_msg = "update could not be completed ('{}')".format(
+                                re.sub(r"\W+", " ", stderr).strip())
+                logger.error(err_msg)
+
+    if not success:
+        if IS_WIN:
+            info_msg = ("for Windows platform it's recommended to use a Github"
+                        "for Windows client for updating purposes "
+                        "(http://windows.github.com/) or just download the "
+                        "latest snapshot from {}".format(SITE))
+        else:
+            info_msg = ("for Linux platform it's required to install a "
+                        "standard 'git' package (e.g.: "
+                        "'sudo apt-get install git')")
+        logger.info(info_msg)

pocsuite/lib/parse/parser.py

@@ -25,6 +25,9 @@ def parseCmdOptions():
     parser.add_argument("--version", action="version",
                         version=VERSION, help="Show program's version number and exit")
 
+    parser.add_argument("--update", dest="update", action="store_true",
+                        help="Update Pocsuite")
+
     target = parser.add_argument_group('target')
 
     target.add_argument("-u", "--url", dest="url",

pocsuite/pocsuite_cli.py

@@ -31,6 +31,7 @@
 from .lib.core.option import init
 from .lib.core.common import delModule
 from .lib.core.common import getUnicode
+from .lib.core.update import update
 
 
 def main():
@@ -73,6 +74,10 @@ def doNothin(*args, **kw):
         dataToStdout("[!] legal disclaimer: %s\n\n" % LEGAL_DISCLAIMER)
         dataToStdout("[*] starting at %s\n\n" % time.strftime("%X"))
 
+        if argsDict['update']:
+            update()
+            return
+
         if argsDict['dork']:
             z = ZoomEye(paths.POCSUITE_RC_PATH)
             if z.newToken():