Allow passing entire mbox via stdin

Per request, allow passing entire mbox files via stdin, allowing fully
pipe-through operation from something like mutt:

    b4 am -sl -m - -o -

Suggested-by: Peter Zijlstra <peterz@infradead.org>
Signed-off-by: Konstantin Ryabitsev <konstantin@linuxfoundation.org>
Link: https://lore.kernel.org/tools/YFETLu8TKWI2WlSF@hirez.programming.kicks-ass.net

3 files changed, 34 insertions, 16 deletions

diff --git a/b4/__init__.py b/b4/__init__.py
index 7e4295e..43fd078 100644
--- a/b4/__init__.py
+++ b/b4/__init__.py
@@ -1956,20 +1956,19 @@ def get_msgid_from_stdin():
     if not sys.stdin.isatty():
         message = email.message_from_string(sys.stdin.read())
         return message.get('Message-ID', None)
-    logger.error('Error: pipe a message or pass msgid as parameter')
-    sys.exit(1)
+    return None
 
 
-def get_msgid(cmdargs):
+def get_msgid(cmdargs) -> Optional[str]:
     if not cmdargs.msgid:
         logger.debug('Getting Message-ID from stdin')
         msgid = get_msgid_from_stdin()
-        if msgid is None:
-            logger.error('Unable to find a valid message-id in stdin.')
-            sys.exit(1)
     else:
         msgid = cmdargs.msgid
 
+    if msgid is None:
+        return None
+
     msgid = msgid.strip('<>')
     # Handle the case when someone pastes a full URL to the message
     matches = re.search(r'^https?://[^/]+/([^/]+)/([^/]+@[^/]+)', msgid, re.IGNORECASE)
diff --git a/b4/command.py b/b4/command.py
index 5a9d6eb..4c0a788 100644
--- a/b4/command.py
+++ b/b4/command.py
@@ -25,7 +25,7 @@ def cmd_mbox_common_opts(sp):
     sp.add_argument('-n', '--mbox-name', dest='wantname', default=None,
                     help='Filename to name the mbox file')
     sp.add_argument('-m', '--use-local-mbox', dest='localmbox', default=None,
-                    help='Instead of grabbing a thread from lore, process this mbox file')
+                    help='Instead of grabbing a thread from lore, process this mbox file (or - for stdin)')
     sp.add_argument('-C', '--no-cache', dest='nocache', action='store_true', default=False,
                     help='Do not use local cache')
 
diff --git a/b4/mbox.py b/b4/mbox.py
index f0bae35..c9d5715 100644
--- a/b4/mbox.py
+++ b/b4/mbox.py
@@ -506,31 +506,50 @@ def main(cmdargs):
         cmdargs.nocache = True
 
     savefile = mkstemp('b4-mbox')[1]
-    msgid = b4.get_msgid(cmdargs)
+    msgid = None
 
     if not cmdargs.localmbox:
+        msgid = b4.get_msgid(cmdargs)
+        if not msgid:
+            logger.error('Error: pipe a message or pass msgid as parameter')
+            os.unlink(savefile)
+            sys.exit(1)
+
         threadfile = b4.get_pi_thread_by_msgid(msgid, savefile, useproject=cmdargs.useproject, nocache=cmdargs.nocache)
         if threadfile is None:
             os.unlink(savefile)
             return
     else:
-        if os.path.exists(cmdargs.localmbox):
+        if cmdargs.localmbox == '-':
+            # The entire mbox is passed via stdin, so save it into a temporary file
+            # and use the first message for our msgid
+            with open(savefile, 'wb') as fh:
+                fh.write(sys.stdin.buffer.read())
+            in_mbx = mailbox.mbox(savefile)
+            msg = in_mbx.get(0)
+            if msg:
+                msgid = msg.get('Message-ID', None).strip('<>')
+
+        elif os.path.exists(cmdargs.localmbox):
+            msgid = b4.get_msgid(cmdargs)
             if os.path.isdir(cmdargs.localmbox):
                 in_mbx = mailbox.Maildir(cmdargs.localmbox)
             else:
                 in_mbx = mailbox.mbox(cmdargs.localmbox)
-            out_mbx = mailbox.mbox(savefile)
-            b4.save_strict_thread(in_mbx, out_mbx, msgid)
-            if not len(out_mbx):
-                logger.critical('Could not find %s in %s', msgid, cmdargs.localmbox)
-                os.unlink(savefile)
-                sys.exit(1)
-            threadfile = savefile
+            if msgid:
+                out_mbx = mailbox.mbox(savefile)
+                b4.save_strict_thread(in_mbx, out_mbx, msgid)
+                if not len(out_mbx):
+                    logger.critical('Could not find %s in %s', msgid, cmdargs.localmbox)
+                    os.unlink(savefile)
+                    sys.exit(1)
         else:
             logger.critical('Mailbox %s does not exist', cmdargs.localmbox)
             os.unlink(savefile)
             sys.exit(1)
 
+        threadfile = savefile
+
     if threadfile and cmdargs.checknewer:
         get_extra_series(threadfile, direction=1)