diff --git a/kindle-download/kindle-download.rb b/kindle-download/kindle-download.rb
index c78f3b4..c03d0c8 100755
--- a/kindle-download/kindle-download.rb
+++ b/kindle-download/kindle-download.rb
@@ -32,11 +32,16 @@ Book = Data.define(:id, :title)
 AMAZON_EMAIL = ENV.fetch("AMAZON_EMAIL")
 AMAZON_PASSWORD = ENV.fetch("AMAZON_PASSWORD")
 
+def downloaded_books
+  Dir[File.expand_path("~/Downloads/*")]
+    .select { it =~ /\.(azw3?|tpz)$/ } # tpz?!
+end
+
 def wait_for_download
-  already_downloaded = Dir[File.expand_path("~/Downloads/*.azw*")]
+  already_downloaded = downloaded_books
   yield
   loop do
-    current = Dir[File.expand_path("~/Downloads/*.azw*")]
+    current = downloaded_books
     new_book, *rest = current - already_downloaded
     if new_book.nil?
       sleep 0.1
@@ -61,7 +66,7 @@ click_on "mfaSubmit"
 
 # I could use pstore here or something, but this is
 # pretty damn convenient
-files = Dir[File.expand_path("~/Downloads/*.azw*")]
+files = downloaded_books
 ids = files.map {
   `xattr -p com.apple.metadata:kMDItemWhereFroms "#{it}"`.chomp
 }