From cb7526fca0aa6cb07d279fec93288d0169904bfc Mon Sep 17 00:00:00 2001 From: AnotherTest Date: Sat, 29 Aug 2020 15:38:36 +0430 Subject: [PATCH] Userland: Make 'pro' capable of guessing a filename The whole thing with `pro url > filename` was getting annoying, so let's just have it support Content-Disposition and guessing from the URL. --- Userland/pro.cpp | 132 ++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 130 insertions(+), 2 deletions(-) diff --git a/Userland/pro.cpp b/Userland/pro.cpp index 18b8dd4141..4138f5568f 100644 --- a/Userland/pro.cpp +++ b/Userland/pro.cpp @@ -24,20 +24,116 @@ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. */ +#include +#include #include #include #include #include #include +#include #include #include #include +// FIXME: Move this somewhere else when it's needed (e.g. in the Browser) +class ContentDispositionParser { +public: + ContentDispositionParser(const StringView& value) + { + GenericLexer lexer(value); + + lexer.consume_while([](auto c) { return is_whitespace(c); }); + + if (lexer.consume_specific("inline")) { + m_kind = Kind::Inline; + if (!lexer.is_eof()) + m_might_be_wrong = true; + return; + } + + if (lexer.consume_specific("attachment")) { + m_kind = Kind::Attachment; + if (lexer.consume_specific(";")) { + lexer.consume_while([](auto c) { return is_whitespace(c); }); + if (lexer.consume_specific("filename=")) { + m_filename = lexer.consume_quoted_string(); + } else { + m_might_be_wrong = true; + } + } + return; + } + + if (lexer.consume_specific("form-data")) { + m_kind = Kind::FormData; + while (lexer.consume_specific(";")) { + lexer.consume_while([](auto c) { return is_whitespace(c); }); + if (lexer.consume_specific("name=")) { + m_name = lexer.consume_quoted_string(); + } else if (lexer.consume_specific("filename=")) { + m_filename = lexer.consume_quoted_string(); + } else { + m_might_be_wrong = true; + } + } + + return; + } + + // FIXME: Support 'filename*' + m_might_be_wrong = true; + } + + enum class Kind { + Inline, + Attachment, + FormData, + }; + + const StringView& filename() const { return m_filename; } + const StringView& name() const { return m_name; } + Kind kind() const { return m_kind; } + bool might_be_wrong() const { return m_might_be_wrong; } + +private: + StringView m_filename; + StringView m_name; + Kind m_kind { Kind::Inline }; + bool m_might_be_wrong { false }; +}; + +static void do_write(const ByteBuffer& payload) +{ + size_t length_remaining = payload.size(); + size_t length_written = 0; + while (length_remaining > 0) { + auto nwritten = fwrite(payload.offset_pointer(length_written), sizeof(char), length_remaining, stdout); + if (nwritten > 0) { + length_remaining -= nwritten; + length_written += nwritten; + continue; + } + + if (feof(stdout)) { + fprintf(stderr, "pro: unexpected eof while writing\n"); + return; + } + + if (ferror(stdout)) { + fprintf(stderr, "pro: error while writing\n"); + return; + } + } +} + int main(int argc, char** argv) { const char* url_str = nullptr; + bool save_at_provided_name = false; Core::ArgsParser args_parser; + args_parser.add_option(save_at_provided_name, "Write to a file named as the remote file", nullptr, 'O'); args_parser.add_positional_argument(url_str, "URL to download from", "url"); args_parser.parse(argc, argv); @@ -55,6 +151,7 @@ int main(int argc, char** argv) fprintf(stderr, "Failed to start download for '%s'\n", url_str); return 1; } + u32 previous_downloaded_size { 0 }; timeval prev_time, current_time, time_diff; gettimeofday(&prev_time, nullptr); @@ -79,11 +176,42 @@ int main(int argc, char** argv) previous_downloaded_size = downloaded_size; prev_time = current_time; }; - download->on_finish = [&](bool success, auto& payload, auto, auto&, auto) { + download->on_finish = [&](bool success, auto& payload, auto, auto& response_headers, auto) { fprintf(stderr, "\033]9;-1;\033\\"); fprintf(stderr, "\n"); + if (success && save_at_provided_name) { + String output_name; + if (auto content_disposition = response_headers.get("Content-Disposition"); content_disposition.has_value()) { + auto& value = content_disposition.value(); + ContentDispositionParser parser(value); + output_name = parser.filename(); + } + + if (output_name.is_empty()) + output_name = url.path(); + + LexicalPath path { output_name }; + output_name = path.basename(); + + // The URL didn't have a name component, e.g. 'serenityos.org' + if (output_name.is_empty() || output_name == "/") { + int i = -1; + do { + output_name = url.host(); + if (i > -1) + output_name = String::format("%s.%d", output_name.characters(), i); + ++i; + } while (Core::File::exists(output_name)); + } + + if (freopen(output_name.characters(), "w", stdout) == nullptr) { + perror("freopen"); + success = false; // oops! + loop.quit(1); + } + } if (success) - write(STDOUT_FILENO, payload.data(), payload.size()); + do_write(payload); else fprintf(stderr, "Download failed :(\n"); loop.quit(0);