From b213f7d999ce9958eea452a468cebe3fc4bd24aa Mon Sep 17 00:00:00 2001
From: WINEZERO <kissloo@foxmail.com>
Date: Wed, 10 May 2023 10:40:48 +0800
Subject: [PATCH] =?UTF-8?q?=E6=94=AF=E6=8C=81=E6=A0=B9=E6=8D=AE=E5=93=8D?=
 =?UTF-8?q?=E5=BA=94=E6=9B=B4=E6=96=B0=20cookie?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

合并https://github.com/urbanadventurer/WhatWeb/pull/281

支持根据响应更新 cookie
---
 lib/target.rb          | 15 +++++++++++++++
 lib/whatweb.rb         |  1 +
 lib/whatweb/version.rb |  2 +-
 whatweb.rb             |  4 ++++
 4 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/lib/target.rb b/lib/target.rb
index 6e1a7ab..6455b3f 100644
--- a/lib/target.rb
+++ b/lib/target.rb
@@ -356,6 +356,21 @@ def open_url(options)
       end
 
       @headers['set-cookie'] = res.get_fields('set-cookie').join("\n") unless @headers['set-cookie'].nil?
+      # update cookies
+      if $UPDATE_COOKIES and not @headers['set-cookie'].nil?
+        res.get_fields('set-cookie').each do |raw_cookie|
+          cookie = raw_cookie.split(';')[0]
+          cookie_name = cookie.split('=')[0]
+          cookie_pattern = Regexp.new('%s=.*?(?=$|;)' % cookie_name)
+          if not $CUSTOM_HEADERS.has_key?('Cookie')
+            $CUSTOM_HEADERS['Cookie'] = cookie
+          elsif $CUSTOM_HEADERS['Cookie'].match(cookie_pattern)
+            $CUSTOM_HEADERS['Cookie'].sub!(cookie_pattern, cookie)
+          else
+            $CUSTOM_HEADERS['Cookie'] += '; %s' % cookie
+          end
+        end
+      end
 
       @status = res.code.to_i
       puts @uri.to_s + " [#{status}]" if $verbose > 1
diff --git a/lib/whatweb.rb b/lib/whatweb.rb
index 1e2631d..7f64791 100644
--- a/lib/whatweb.rb
+++ b/lib/whatweb.rb
@@ -67,6 +67,7 @@
 $USER_AGENT = "WhatWeb/#{WhatWeb::VERSION}"
 $AGGRESSION = 1
 $FOLLOW_REDIRECT = 'always'
+$UPDATE_COOKIES = false
 $USE_PROXY = false
 $PROXY_HOST = nil
 $PROXY_PORT = 8080
diff --git a/lib/whatweb/version.rb b/lib/whatweb/version.rb
index ad7a74c..e1e4092 100644
--- a/lib/whatweb/version.rb
+++ b/lib/whatweb/version.rb
@@ -15,5 +15,5 @@
 # You should have received a copy of the GNU General Public License
 # along with WhatWeb.  If not, see <http://www.gnu.org/licenses/>.
 module WhatWeb
-  VERSION = '0.5.5.14'.freeze
+  VERSION = '0.5.5.15'.freeze
 end
diff --git a/whatweb.rb b/whatweb.rb
index 789986e..32f7665 100644
--- a/whatweb.rb
+++ b/whatweb.rb
@@ -92,6 +92,7 @@ def usage_full
   --user, -u=<user:password>\tHTTP basic authentication.
   --cookie, -c=COOKIES\t\tUse cookies, e.g. 'name=value; name2=value2'.
   --cookie-jar=FILE\t\tRead cookies from a file.
+  --update-cookies\t\tUpdate cookies according to responses.
 
 PROXY:
   --proxy, -P\t\t\t<hostname[:port]> Set proxy hostname and port.
@@ -347,6 +348,7 @@ def error(s)
   ['--header', '-H', GetoptLong::REQUIRED_ARGUMENT],
   ['--cookie', '-c', GetoptLong::REQUIRED_ARGUMENT],
   ['--cookie-jar', GetoptLong::REQUIRED_ARGUMENT],
+  ['--update-cookies', GetoptLong::NO_ARGUMENT],
   ['--user', '-u', GetoptLong::REQUIRED_ARGUMENT],
   ['--wait', GetoptLong::REQUIRED_ARGUMENT],
   ['--debug', GetoptLong::NO_ARGUMENT],
@@ -530,6 +532,8 @@ def error(s)
       rescue
         raise("Could not read cookies from #{arg}.")
       end
+    when '--update-cookies'
+      $UPDATE_COOKIES = true
     when '-u', '--user'
       unless arg.include?(':')
         raise("Incorrect credentials format.")