import/export daemon: Move command building into class

Instead of passing around many variables for building the executed
command, they're now kept as instance variables.

Signed-off-by: Michael Hanselmann <hansmi@google.com>
Reviewed-by: Guido Trotter <ultrotter@google.com>

diff --git a/daemons/import-export b/daemons/import-export
index 134b3e9..7e9b265 100755 (executable)
--- a/daemons/import-export
+++ b/daemons/import-export
@@ -291,138 +291,148 @@ def ProcessOutput(line, status_file, logger, socat):
   status_file.Update(force_update)
 
 
-def GetBashCommand(cmd):
-  """Prepares a command to be run in Bash.
+class CommandBuilder(object):
+  def __init__(self, mode, opts, socat_stderr_fd):
+    """Initializes this class.
 
-  """
-  return ["bash", "-o", "errexit", "-o", "pipefail", "-c", cmd]
+    @param mode: Daemon mode (import or export)
+    @param opts: Options object
+    @type socat_stderr_fd: int
+    @param socat_stderr_fd: File descriptor socat should write its stderr to
 
+    """
+    self._opts = opts
+    self._mode = mode
+    self._socat_stderr_fd = socat_stderr_fd
 
-def GetSocatCommand(mode):
-  """Returns the socat command.
+  @staticmethod
+  def GetBashCommand(cmd):
+    """Prepares a command to be run in Bash.
 
-  """
-  common_addr_opts = SOCAT_TCP_OPTS + SOCAT_OPENSSL_OPTS + [
-    "key=%s" % options.key,
-    "cert=%s" % options.cert,
-    "cafile=%s" % options.ca,
-    ]
-
-  if options.bind is not None:
-    common_addr_opts.append("bind=%s" % options.bind)
-
-  if mode == constants.IEM_IMPORT:
-    if options.port is None:
-      port = 0
-    else:
-      port = options.port
+    """
+    return ["bash", "-o", "errexit", "-o", "pipefail", "-c", cmd]
 
-    addr1 = [
-      "OPENSSL-LISTEN:%s" % port,
-      "reuseaddr",
+  def _GetSocatCommand(self):
+    """Returns the socat command.
 
-      # Retry to listen if connection wasn't established successfully, up to
-      # 100 times a second. Note that this still leaves room for DoS attacks.
-      "forever",
-      "intervall=0.01",
-      ] + common_addr_opts
-    addr2 = ["stdout"]
+    """
+    common_addr_opts = SOCAT_TCP_OPTS + SOCAT_OPENSSL_OPTS + [
+      "key=%s" % self._opts.key,
+      "cert=%s" % self._opts.cert,
+      "cafile=%s" % self._opts.ca,
+      ]
+
+    if self._opts.bind is not None:
+      common_addr_opts.append("bind=%s" % self._opts.bind)
+
+    if self._mode == constants.IEM_IMPORT:
+      if self._opts.port is None:
+        port = 0
+      else:
+        port = self._opts.port
 
-  elif mode == constants.IEM_EXPORT:
-    addr1 = ["stdin"]
-    addr2 = [
-      "OPENSSL:%s:%s" % (options.host, options.port),
+      addr1 = [
+        "OPENSSL-LISTEN:%s" % port,
+        "reuseaddr",
 
-      # How long to wait per connection attempt
-      "connect-timeout=%s" % options.connect_timeout,
+        # Retry to listen if connection wasn't established successfully, up to
+        # 100 times a second. Note that this still leaves room for DoS attacks.
+        "forever",
+        "intervall=0.01",
+        ] + common_addr_opts
+      addr2 = ["stdout"]
 
-      # Retry a few times before giving up to connect (once per second)
-      "retry=%s" % options.connect_retries,
-      "intervall=1",
-      ] + common_addr_opts
+    elif self._mode == constants.IEM_EXPORT:
+      addr1 = ["stdin"]
+      addr2 = [
+        "OPENSSL:%s:%s" % (self._opts.host, self._opts.port),
 
-  else:
-    raise Error("Invalid mode")
+        # How long to wait per connection attempt
+        "connect-timeout=%s" % self._opts.connect_timeout,
 
-  for i in [addr1, addr2]:
-    for value in i:
-      if "," in value:
-        raise Error("Comma not allowed in socat option value: %r" % value)
+        # Retry a few times before giving up to connect (once per second)
+        "retry=%s" % self._opts.connect_retries,
+        "intervall=1",
+        ] + common_addr_opts
 
-  return [
-    constants.SOCAT_PATH,
+    else:
+      raise Error("Invalid mode '%s'" % self._mode)
 
-    # Log to stderr
-    "-ls",
+    for i in [addr1, addr2]:
+      for value in i:
+        if "," in value:
+          raise Error("Comma not allowed in socat option value: %r" % value)
 
-    # Log level
-    "-d", "-d",
+    return [
+      constants.SOCAT_PATH,
 
-    # Buffer size
-    "-b%s" % SOCAT_BUFSIZE,
+      # Log to stderr
+      "-ls",
 
-    # Unidirectional mode, the first address is only used for reading, and the
-    # second address is only used for writing
-    "-u",
+      # Log level
+      "-d", "-d",
 
-    ",".join(addr1), ",".join(addr2)
-    ]
+      # Buffer size
+      "-b%s" % SOCAT_BUFSIZE,
 
+      # Unidirectional mode, the first address is only used for reading, and the
+      # second address is only used for writing
+      "-u",
 
-def GetTransportCommand(mode, socat_stderr_fd):
-  """Returns the command for the transport part of the daemon.
+      ",".join(addr1), ",".join(addr2)
+      ]
 
-  @param mode: Daemon mode (import or export)
-  @type socat_stderr_fd: int
-  @param socat_stderr_fd: File descriptor socat should write its stderr to
+  def _GetTransportCommand(self):
+    """Returns the command for the transport part of the daemon.
 
-  """
-  socat_cmd = ("%s 2>&%d" %
-               (utils.ShellQuoteArgs(GetSocatCommand(mode)),
-                socat_stderr_fd))
+    """
+    socat_cmd = ("%s 2>&%d" %
+                 (utils.ShellQuoteArgs(self._GetSocatCommand()),
+                  self._socat_stderr_fd))
 
-  compr = options.compress
+    compr = self._opts.compress
 
-  assert compr in constants.IEC_ALL
+    assert compr in constants.IEC_ALL
 
-  if mode == constants.IEM_IMPORT:
-    if compr == constants.IEC_GZIP:
-      transport_cmd = "%s | gunzip -c" % socat_cmd
-    else:
-      transport_cmd = socat_cmd
-  elif mode == constants.IEM_EXPORT:
-    if compr == constants.IEC_GZIP:
-      transport_cmd = "gzip -c | %s" % socat_cmd
+    if self._mode == constants.IEM_IMPORT:
+      if compr == constants.IEC_GZIP:
+        transport_cmd = "%s | gunzip -c" % socat_cmd
+      else:
+        transport_cmd = socat_cmd
+    elif self._mode == constants.IEM_EXPORT:
+      if compr == constants.IEC_GZIP:
+        transport_cmd = "gzip -c | %s" % socat_cmd
+      else:
+        transport_cmd = socat_cmd
     else:
-      transport_cmd = socat_cmd
-  else:
-    raise Error("Invalid mode")
+      raise Error("Invalid mode '%s'" % self._mode)
 
-  # TODO: Use "dd" to measure processed data (allows to give an ETA)
+    # TODO: Use "dd" to measure processed data (allows to give an ETA)
 
-  # TODO: Run transport as separate user
-  # The transport uses its own shell to simplify running it as a separate user
-  # in the future.
-  return GetBashCommand(transport_cmd)
+    # TODO: Run transport as separate user
+    # The transport uses its own shell to simplify running it as a separate user
+    # in the future.
+    return self.GetBashCommand(transport_cmd)
 
+  def GetCommand(self):
+    """Returns the complete child process command.
 
-def GetCommand(mode, socat_stderr_fd):
-  """Returns the complete child process command.
+    """
+    transport_cmd = self._GetTransportCommand()
 
-  """
-  buf = StringIO()
+    buf = StringIO()
 
-  if options.cmd_prefix:
-    buf.write(options.cmd_prefix)
-    buf.write(" ")
+    if self._opts.cmd_prefix:
+      buf.write(self._opts.cmd_prefix)
+      buf.write(" ")
 
-  buf.write(utils.ShellQuoteArgs(GetTransportCommand(mode, socat_stderr_fd)))
+    buf.write(utils.ShellQuoteArgs(transport_cmd))
 
-  if options.cmd_suffix:
-    buf.write(" ")
-    buf.write(options.cmd_suffix)
+    if self._opts.cmd_suffix:
+      buf.write(" ")
+      buf.write(self._opts.cmd_suffix)
 
-  return GetBashCommand(buf.getvalue())
+    return self.GetBashCommand(buf.getvalue())
 
 
 def ProcessChildIO(child, socat_stderr_read_fd, status_file, child_logger,
@@ -663,7 +673,7 @@ def main():
       (socat_stderr_read_fd, socat_stderr_write_fd) = os.pipe()
 
       # Get child process command
-      cmd = GetCommand(mode, socat_stderr_write_fd)
+      cmd = CommandBuilder(mode, options, socat_stderr_write_fd).GetCommand()
 
       logging.debug("Starting command %r", cmd)