|
| 1 | +__author__ = 'Ondrej Mular <[email protected]>' |
| 2 | +__all__ = ["FencingPyCurl"] |
| 3 | + |
| 4 | +import pycurl |
| 5 | +import sys |
| 6 | +import atexit |
| 7 | +import json |
| 8 | +import StringIO |
| 9 | +import time |
| 10 | +import logging |
| 11 | +import pprint |
| 12 | + |
| 13 | +## do not add code here. |
| 14 | +#BEGIN_VERSION_GENERATION |
| 15 | +RELEASE_VERSION = "" |
| 16 | +REDHAT_COPYRIGHT = "" |
| 17 | +BUILD_DATE = "" |
| 18 | +#END_VERSION_GENERATION |
| 19 | + |
| 20 | +class FencingPyCurl(): |
| 21 | + active = False |
| 22 | + input_file = None |
| 23 | + output_file = None |
| 24 | + options = None |
| 25 | + actions = None |
| 26 | + request_index = 0 |
| 27 | + output_buffer = None |
| 28 | + write_function = None |
| 29 | + pycurl_obj = None |
| 30 | + |
| 31 | + def __init__(self): |
| 32 | + if not FencingPyCurl.active and (FencingPyCurl.input_file or FencingPyCurl.output_file): |
| 33 | + FencingPyCurl.active = True |
| 34 | + logging.debug("FencingPyCurl is active now") |
| 35 | + if FencingPyCurl.active: |
| 36 | + self.options = { |
| 37 | + "request": {}, |
| 38 | + "response": {}, |
| 39 | + "time": {} |
| 40 | + } |
| 41 | + |
| 42 | + self.output_buffer = StringIO.StringIO() |
| 43 | + if FencingPyCurl.actions is None: |
| 44 | + if FencingPyCurl.input_file: |
| 45 | + logging.debug("Reading input file.") |
| 46 | + |
| 47 | + try: |
| 48 | + with open(FencingPyCurl.input_file, "r") as f: |
| 49 | + FencingPyCurl.actions = json.load(f) |
| 50 | + except Exception as e: |
| 51 | + logging.debug("Reading input file (%s) failed: %s" % (FencingPyCurl.input_file, e.message)) |
| 52 | + |
| 53 | + if FencingPyCurl.output_file: |
| 54 | + logging.debug("output file detected") |
| 55 | + if not FencingPyCurl.actions: |
| 56 | + FencingPyCurl.actions = [] |
| 57 | + FencingPyCurl.request_index = 0 |
| 58 | + |
| 59 | + self.pycurl_obj = pycurl.Curl() |
| 60 | + |
| 61 | + def setopt(self, opt, value): |
| 62 | + if FencingPyCurl.active: |
| 63 | + if opt == pycurl.WRITEFUNCTION: |
| 64 | + self.write_function = value |
| 65 | + value = self.output_buffer.write |
| 66 | + else: |
| 67 | + self.options["request"][str(opt)] = value |
| 68 | + return self.pycurl_obj.setopt(opt, value) |
| 69 | + |
| 70 | + def perform(self): |
| 71 | + if FencingPyCurl.active: |
| 72 | + if FencingPyCurl.input_file: |
| 73 | + perform_start = time.time() |
| 74 | + if self.options["request"] == FencingPyCurl.actions[FencingPyCurl.request_index]["request"]: |
| 75 | + self.options["response"] = FencingPyCurl.actions[FencingPyCurl.request_index]["response"] |
| 76 | + if self.write_function: |
| 77 | + self.write_function(self.options["response"]["output"]) |
| 78 | + diff = FencingPyCurl.actions[FencingPyCurl.request_index]["time"]["perform_duration"] - (time.time() - perform_start) |
| 79 | + if diff > 0: |
| 80 | + logging.debug("sleeping for: %s" % str(diff)) |
| 81 | + time.sleep(diff) |
| 82 | + FencingPyCurl.last_request_time = time.time() |
| 83 | + else: |
| 84 | + print "Request:" |
| 85 | + pprint.pprint(self.options["request"]) |
| 86 | + print "Expected:" |
| 87 | + pprint.pprint(FencingPyCurl.actions[FencingPyCurl.request_index]["request"]) |
| 88 | + raise Exception("Invalid request") |
| 89 | + else: |
| 90 | + start_time = time.time() |
| 91 | + self.pycurl_obj.perform() |
| 92 | + self.options["time"]["perform_duration"] = FencingPyCurl.last_request_time - start_time |
| 93 | + self.options["response"]["output"] = self.output_buffer.getvalue() |
| 94 | + if self.write_function: |
| 95 | + self.write_function(self.options["response"]["output"]) |
| 96 | + if FencingPyCurl.output_file: |
| 97 | + FencingPyCurl.actions.append(self.options) |
| 98 | + FencingPyCurl.request_index += 1 |
| 99 | + else: |
| 100 | + return self.pycurl_obj.perform() |
| 101 | + |
| 102 | + def unsetopt(self, opt): |
| 103 | + if FencingPyCurl.active: |
| 104 | + del self.options["request"][str(opt)] |
| 105 | + return self.pycurl_obj.unsetopt(opt) |
| 106 | + |
| 107 | + def getinfo(self, opt): |
| 108 | + value = self.pycurl_obj.getinfo(opt) |
| 109 | + if FencingPyCurl.active: |
| 110 | + if FencingPyCurl.input_file: |
| 111 | + value = self.options["response"][str(opt)] |
| 112 | + else: |
| 113 | + self.options["response"][opt] = value |
| 114 | + return value |
| 115 | + |
| 116 | + def reset(self): |
| 117 | + if FencingPyCurl.active: |
| 118 | + self.options.clear() |
| 119 | + return self.pycurl_obj.reset() |
| 120 | + |
| 121 | + def close(self): |
| 122 | + return self.pycurl_obj.close() |
| 123 | + |
| 124 | + @staticmethod |
| 125 | + def save_log_to_file(): |
| 126 | + if FencingPyCurl.output_file and FencingPyCurl.actions: |
| 127 | + logging.debug("Writing log to file: %s" % FencingPyCurl.output_file) |
| 128 | + try: |
| 129 | + with open(FencingPyCurl.output_file, "w") as f: |
| 130 | + json.dump(FencingPyCurl.actions, f, sort_keys=True, indent=4, separators=(',', ': ')) |
| 131 | + except Exception as e: |
| 132 | + logging.debug("Writing log to file (%s) failed: %s" % (FencingPyCurl.output_file, e.message)) |
| 133 | + |
| 134 | + |
| 135 | +def get_and_remove_arg(arg, has_value=True): |
| 136 | + logging.debug("Getting arg: %s (has_value: %s)" % (arg, str(has_value))) |
| 137 | + if not has_value: |
| 138 | + if arg in sys.argv: |
| 139 | + sys.argv.remove(arg) |
| 140 | + logging.debug("%s: True" % arg) |
| 141 | + return True |
| 142 | + if arg in sys.argv: |
| 143 | + index = sys.argv.index(arg) |
| 144 | + sys.argv.remove(arg) |
| 145 | + if len(sys.argv) > index: |
| 146 | + value = sys.argv[index] |
| 147 | + sys.argv.remove(value) |
| 148 | + logging.debug("%s: %s" % (arg, value)) |
| 149 | + return value |
| 150 | + return None |
| 151 | + |
| 152 | +FencingPyCurl.input_file = get_and_remove_arg("--fencing_pycurl-log-in") |
| 153 | +FencingPyCurl.output_file = get_and_remove_arg("--fencing_pycurl-log-out") |
| 154 | + |
| 155 | +if FencingPyCurl.output_file: |
| 156 | + atexit.register(FencingPyCurl.save_log_to_file) |
0 commit comments