diff --git a/lib/paygen/filelib.py b/lib/paygen/filelib.py
index 18e5f86..5f5ffe6 100644
--- a/lib/paygen/filelib.py
+++ b/lib/paygen/filelib.py
@@ -99,3 +99,24 @@
   sha256_hex = base64.b64encode(sha256.digest())
 
   return sha1_hex, sha256_hex
+
+
+def CopyFileSegment(in_file, in_mode, in_len, out_file, out_mode, in_seek=0):
+  """Simulates a `dd` operation with seeks.
+
+  Args:
+    in_file: The input file
+    in_mode: The mode to open the input file
+    in_len: The length to copy
+    out_file: The output file
+    out_mode: The mode to open the output file
+    in_seek: How many bytes to seek from the |in_file|
+  """
+  with open(in_file, in_mode) as in_stream, \
+       open(out_file, out_mode) as out_stream:
+    in_stream.seek(in_seek)
+    remaining = in_len
+    while remaining:
+      chunk = in_stream.read(min(8192 * 1024, remaining))
+      remaining -= len(chunk)
+      out_stream.write(chunk)
diff --git a/lib/paygen/filelib_unittest.py b/lib/paygen/filelib_unittest.py
index 17cf5ba..76acd66 100644
--- a/lib/paygen/filelib_unittest.py
+++ b/lib/paygen/filelib_unittest.py
@@ -89,3 +89,12 @@
 
     filelib.Copy(path1, relative_path)
     self.assertExists(path2)
+
+  def testCopyFileSegment(self):
+    """Test copying on a simple file segment."""
+    a = os.path.join(self.tempdir, 'a.txt')
+    osutils.WriteFile(a, '789')
+    b = os.path.join(self.tempdir, 'b.txt')
+    osutils.WriteFile(b, '0123')
+    filelib.CopyFileSegment(a, 'r', 2, b, 'r+', in_seek=1)
+    self.assertEqual(osutils.ReadFile(b), '8923')
diff --git a/lib/paygen/partition_lib.py b/lib/paygen/partition_lib.py
new file mode 100644
index 0000000..dd9aa2c
--- /dev/null
+++ b/lib/paygen/partition_lib.py
@@ -0,0 +1,116 @@
+# -*- coding: utf-8 -*-
+# Copyright 2018 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Library for handling Chrome OS partition."""
+
+from __future__ import print_function
+
+import os
+import shutil
+import tempfile
+
+from chromite.lib import constants
+from chromite.lib import cros_build_lib
+from chromite.lib import cros_logging as logging
+
+from chromite.lib.paygen import filelib
+
+
+def ExtractPartition(filename, partition, out_part):
+  """Extracts partition from an image file.
+
+  Args:
+    filename: The image file.
+    partition: The partition name. e.g. ROOT or KERNEL.
+    out_part: The output partition file.
+  """
+  parts = cros_build_lib.GetImageDiskPartitionInfo(filename)
+  part_info = [p for p in parts if p.name == partition][0]
+
+  offset = int(part_info.start)
+  length = int(part_info.size)
+
+  filelib.CopyFileSegment(filename, 'r', length, out_part, 'w', in_seek=offset)
+
+
+def Ext2FileSystemSize(ext2_file):
+  """Return the size of an ext2 filesystem in bytes.
+
+  Args:
+    ext2_file: The path to the ext2 file.
+  """
+  # dumpe2fs is normally installed in /sbin but doesn't require root.
+  dump = cros_build_lib.RunCommand(['/sbin/dumpe2fs', '-h', ext2_file],
+                                   print_cmd=False, capture_output=True).output
+  fs_blocks = 0
+  fs_blocksize = 0
+  for line in dump.split('\n'):
+    if not line:
+      continue
+    label, data = line.split(':')[:2]
+    if label == 'Block count':
+      fs_blocks = int(data)
+    elif label == 'Block size':
+      fs_blocksize = int(data)
+
+  return fs_blocks * fs_blocksize
+
+
+def PatchKernel(image, kern_file):
+  """Patches a kernel with vblock from a stateful partition.
+
+  Args:
+    image: The stateful partition image.
+    kern_file: The kernel file.
+  """
+
+  with tempfile.NamedTemporaryFile(prefix='stateful') as state_out, \
+       tempfile.NamedTemporaryFile(prefix='vmlinuz_hd.vblock') as vblock:
+    ExtractPartition(image, constants.PART_STATE, state_out)
+    cros_build_lib.RunCommand(
+        ['e2cp', '%s:/vmlinuz_hd.vblock' % state_out, vblock])
+    filelib.CopyFileSegment(
+        vblock, 'r', os.path.getsize(vblock), kern_file, 'r+')
+
+
+def ExtractKernel(image, kern_out):
+  """Extracts the kernel from the given image.
+
+  Args:
+    image: The image containing the kernel partition.
+    kern_out: The output kernel file.
+  """
+  ExtractPartition(image, constants.PART_KERN_B, kern_out)
+  with open(kern_out, 'r') as kern:
+    if not any(kern.read(65536)):
+      logging.warn('%s: Kernel B is empty, patching kernel A.', image)
+      ExtractPartition(image, constants.PART_KERN_A, kern_out)
+      PatchKernel(image, kern_out)
+
+
+def ExtractRoot(image, root_out, root_pretruncate=None):
+  """Extract the rootfs partition from a gpt image.
+
+  Args:
+    image: The input image file.
+    root_out: The output root partition file.
+    root_pretruncate: The output root partition before being truncated.
+  """
+  ExtractPartition(image, constants.PART_ROOT_A, root_out)
+
+  if root_pretruncate:
+    logging.info('Saving pre-truncated root as %s.', root_pretruncate)
+    shutil.copyfile(root_out, root_pretruncate)
+
+  # We only update the filesystem part of the partition, which is stored in the
+  # gpt script.
+  root_out_size = Ext2FileSystemSize(root_out)
+  if root_out_size:
+    with open(root_out, 'a') as root:
+      logging.info('Root size currently %d bytes.', os.path.getsize(root_out))
+      root.truncate(root_out_size)
+    logging.info('Truncated root to %d bytes.', root_out_size)
+  else:
+    raise IOError('Error truncating the rootfs to filesystem size.')
diff --git a/lib/paygen/partition_lib_unittest b/lib/paygen/partition_lib_unittest
new file mode 120000
index 0000000..ef3e37b
--- /dev/null
+++ b/lib/paygen/partition_lib_unittest
@@ -0,0 +1 @@
+../../scripts/wrapper.py
\ No newline at end of file
diff --git a/lib/paygen/partition_lib_unittest.py b/lib/paygen/partition_lib_unittest.py
new file mode 100644
index 0000000..3b49558
--- /dev/null
+++ b/lib/paygen/partition_lib_unittest.py
@@ -0,0 +1,76 @@
+# -*- coding: utf-8 -*-
+# Copyright 2018 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Test the partition_lib module."""
+
+from __future__ import print_function
+
+import os
+
+from chromite.lib import cros_build_lib
+from chromite.lib import cros_test_lib
+from chromite.lib import osutils
+
+from chromite.lib.paygen import partition_lib
+
+class ExtractRootTest(cros_test_lib.MockTempDirTestCase):
+  """Test partition_lib.ExtractRoot"""
+  def testTruncate(self):
+    """Test truncating on extraction."""
+    root = os.path.join(self.tempdir, 'root.bin')
+    root_pretruncate = os.path.join(self.tempdir, 'root_pretruncate.bin')
+
+    content = '0123456789'
+    osutils.WriteFile(root, content)
+    self.PatchObject(partition_lib, 'ExtractPartition',
+                     return_value=root)
+    self.PatchObject(partition_lib, 'Ext2FileSystemSize',
+                     return_value=2)
+
+    partition_lib.ExtractRoot(None, root, root_pretruncate)
+    self.assertEqual(osutils.ReadFile(root), content[:2])
+    self.assertEqual(osutils.ReadFile(root_pretruncate), content)
+
+
+class Ext2FileSystemSizeTest(cros_test_lib.MockTestCase):
+  """Test partition_lib.Ext2FileSystemSize"""
+  def testExt2FileSystemSize(self):
+    """Test on simple output."""
+    block_size = 4096
+    block_count = 123
+
+    self.PatchObject(cros_build_lib, 'RunCommand',
+                     return_value=cros_build_lib.CommandResult(output='''
+Block size: %d
+Other thing: 123456798
+Not an integer: cdsad132csda
+Block count: %d
+''' % (block_size, block_count)))
+
+    size = partition_lib.Ext2FileSystemSize('/dev/null')
+    self.assertEqual(size, block_size * block_count)
+
+
+class ExtractPartitionTest(cros_test_lib.MockTempDirTestCase):
+  """Test partition_lib.ExtractPartition"""
+  def testExtractPartition(self):
+    """Tests extraction on a simple image."""
+    part_a_bin = '0123'
+    part_b_bin = '4567'
+    image_bin = part_a_bin + part_b_bin
+
+    image = os.path.join(self.tempdir, 'image.bin')
+    osutils.WriteFile(image, image_bin)
+    part_a = os.path.join(self.tempdir, 'a.bin')
+
+    fake_partitions = (
+        cros_build_lib.PartitionInfo(1, 0, 4, 4, 'fs', 'PART-A', ''),
+        cros_build_lib.PartitionInfo(2, 4, 8, 4, 'fs', 'PART-B', ''),
+    )
+    self.PatchObject(cros_build_lib, 'GetImageDiskPartitionInfo',
+                     return_value=fake_partitions)
+
+    partition_lib.ExtractPartition(image, 'PART-A', part_a)
+    self.assertEqual(osutils.ReadFile(part_a), part_a_bin)
diff --git a/scripts/cros_generate_update_payload.py b/scripts/cros_generate_update_payload.py
index aeb0431..a78f479 100644
--- a/scripts/cros_generate_update_payload.py
+++ b/scripts/cros_generate_update_payload.py
@@ -10,7 +10,6 @@
 
 from __future__ import print_function
 
-import os
 import shutil
 import tempfile
 
@@ -20,118 +19,12 @@
 from chromite.lib import cros_logging as logging
 from chromite.lib import osutils
 
+from chromite.lib.paygen import partition_lib
+
 
 _DELTA_GENERATOR = 'delta_generator'
 
 
-# TODO(tbrindus): move this to paygen/filelib.py.
-def CopyFileSegment(in_file, in_mode, in_len, out_file, out_mode, in_seek=0):
-  """Simulates a `dd` operation with seeks."""
-  with open(in_file, in_mode) as in_stream, \
-       open(out_file, out_mode) as out_stream:
-    in_stream.seek(in_seek)
-    remaining = in_len
-    while remaining:
-      chunk = in_stream.read(min(8192 * 1024, remaining))
-      remaining -= len(chunk)
-      out_stream.write(chunk)
-
-
-# TODO(tbrindus): move this to paygen/filelib.py.
-def ExtractPartitionToTempFile(filename, partition, temp_file=None):
-  """Extracts |partition| from |filename| into |temp_file|.
-
-  If |temp_file| is not specified, an arbitrary file is used.
-
-  Returns the location of the extracted partition.
-  """
-  if temp_file is None:
-    temp_file = tempfile.mktemp(prefix='cros_generate_update_payload')
-
-  parts = cros_build_lib.GetImageDiskPartitionInfo(filename)
-  part_info = [p for p in parts if p.name == partition][0]
-
-  offset = int(part_info.start)
-  length = int(part_info.size)
-
-  CopyFileSegment(filename, 'r', length, temp_file, 'w', in_seek=offset)
-
-  return temp_file
-
-
-def PatchKernel(image, kern_file):
-  """Patches kernel |kern_file| with vblock from |image| stateful partition."""
-  state_out = ExtractPartitionToTempFile(image, constants.PART_STATE)
-  vblock = tempfile.mktemp(prefix='vmlinuz_hd.vblock')
-  cros_build_lib.RunCommand(['e2cp', '%s:/vmlinuz_hd.vblock' % state_out,
-                             vblock])
-
-  CopyFileSegment(vblock, 'r', os.path.getsize(vblock), kern_file, 'r+')
-
-  osutils.SafeUnlink(state_out)
-  osutils.SafeUnlink(vblock)
-
-
-def ExtractKernel(bin_file, kern_out):
-  """Extracts the kernel from the given |bin_file|, into |kern_out|."""
-  kern_out = ExtractPartitionToTempFile(bin_file, constants.PART_KERN_B,
-                                        kern_out)
-  with open(kern_out, 'r') as kern:
-    if not any(kern.read(65536)):
-      logging.warn('%s: Kernel B is empty, patching kernel A.', bin_file)
-      ExtractPartitionToTempFile(bin_file, constants.PART_KERN_A, kern_out)
-      PatchKernel(bin_file, kern_out)
-
-  return kern_out
-
-
-def Ext2FileSystemSize(rootfs):
-  """Return the size in bytes of the ext2 filesystem passed in |rootfs|."""
-  # dumpe2fs is normally installed in /sbin but doesn't require root.
-  dump = cros_build_lib.RunCommand(['/sbin/dumpe2fs', '-h', rootfs],
-                                   print_cmd=False, capture_output=True).output
-  fs_blocks = 0
-  fs_blocksize = 0
-  for line in dump.split('\n'):
-    if not line:
-      continue
-    label, data = line.split(':')[:2]
-    if label == 'Block count':
-      fs_blocks = int(data)
-    elif label == 'Block size':
-      fs_blocksize = int(data)
-
-  return fs_blocks * fs_blocksize
-
-
-def ExtractRoot(bin_file, root_out, root_pretruncate=None):
-  """Extract the rootfs partition from the gpt image |bin_file|.
-
-  Stores it in |root_out|. If |root_out| is empty, a new temp file will be used.
-  If |root_pretruncate| is non-empty, saves the pretruncated rootfs partition
-  there.
-  """
-  root_out = ExtractPartitionToTempFile(bin_file, constants.PART_ROOT_A,
-                                        root_out)
-
-  if root_pretruncate:
-    logging.info('Saving pre-truncated root as %s.', root_pretruncate)
-    shutil.copyfile(root_out, root_pretruncate)
-
-  # We only update the filesystem part of the partition, which is stored in the
-  # gpt script.
-  root_out_size = Ext2FileSystemSize(root_out)
-  if root_out_size:
-    with open(root_out, 'a') as root:
-      logging.info('Root size currently %d bytes.', os.path.getsize(root_out))
-      root.truncate(root_out_size)
-    logging.info('Truncated root to %d bytes.', root_out_size)
-  else:
-    raise IOError('Error truncating the rootfs to filesystem size.')
-
-  return root_out
-
-
 def GenerateUpdatePayload(opts):
   """Generates the output files for the given commandline |opts|."""
   # TODO(tbrindus): we can support calling outside of chroot easily with
@@ -150,13 +43,14 @@
       if opts.src_image:
         src_kernel_path = opts.src_kern_path or 'old_kern.dat'
         src_root_path = opts.src_root_path or 'old_root.dat'
-        ExtractKernel(opts.src_image, src_kernel_path)
-        ExtractRoot(opts.src_image, src_root_path)
+        partition_lib.ExtractKernel(opts.src_image, src_kernel_path)
+        partition_lib.ExtractRoot(opts.src_image, src_root_path)
       if opts.image:
         dst_kernel_path = opts.kern_path or 'new_kern.dat'
         dst_root_path = opts.root_path or 'new_root.dat'
-        ExtractKernel(opts.image, dst_kernel_path)
-        ExtractRoot(opts.image, dst_root_path, opts.root_pretruncate_path)
+        partition_lib.ExtractKernel(opts.image, dst_kernel_path)
+        partition_lib.ExtractRoot(opts.image, dst_root_path,
+                                  opts.root_pretruncate_path)
       logging.info('Done extracting kernel/root')
       return
 
@@ -164,17 +58,24 @@
 
     logging.info('Generating %s update', payload_type)
 
+    src_kernel_path = ''
+    src_root_path = ''
     if delta:
       if not opts.full_kernel:
-        src_kernel_path = ExtractKernel(opts.src_image, opts.src_kern_path)
+        src_kernel_path = (opts.src_kern_path or
+                           tempfile.NamedTemporaryFile().name)
+        partition_lib.ExtractKernel(opts.src_image, src_kernel_path)
       else:
         logging.info('Generating full kernel update')
 
-      src_root_path = ExtractRoot(opts.src_image, opts.src_root_path)
+      src_root_path = opts.src_root_path or tempfile.NamedTemporaryFile().name
+      partition_lib.ExtractRoot(opts.src_image, src_root_path)
 
-    dst_kernel_path = ExtractKernel(opts.image, opts.kern_path)
-    dst_root_path = ExtractRoot(opts.image, opts.root_path,
-                                opts.root_pretruncate_path)
+    dst_kernel_path = opts.kern_path or tempfile.NamedTemporaryFile().name
+    dst_root_path = opts.root_path or tempfile.NamedTemporaryFile().name
+    partition_lib.ExtractKernel(opts.image, dst_kernel_path)
+    partition_lib.ExtractRoot(opts.image, dst_root_path,
+                              opts.root_pretruncate_path)
 
     # In major version 2 we need to explicity mark the postinst on the root
     # partition to run.
diff --git a/scripts/cros_generate_update_payload_unittest.py b/scripts/cros_generate_update_payload_unittest.py
index a3b935a..219af01 100644
--- a/scripts/cros_generate_update_payload_unittest.py
+++ b/scripts/cros_generate_update_payload_unittest.py
@@ -13,98 +13,17 @@
 from chromite.lib import cros_build_lib
 from chromite.lib import cros_test_lib
 from chromite.lib import osutils
+from chromite.lib.paygen import partition_lib
 from chromite.scripts import cros_generate_update_payload
 
-
-class CopyFileSegmentTest(cros_test_lib.TempDirTestCase):
-  """Test cros_generate_update_payload.CopyFileSegment"""
-  def testCopyFileSegment(self):
-    """Test copying on a simple file."""
-    a = os.path.join(self.tempdir, 'a.txt')
-    osutils.WriteFile(a, '789')
-    b = os.path.join(self.tempdir, 'b.txt')
-    osutils.WriteFile(b, '0123')
-    cros_generate_update_payload.CopyFileSegment(a, 'r', 2, b, 'r+', in_seek=1)
-    self.assertEqual(osutils.ReadFile(b), '8923')
-
-
-class ExtractRootTest(cros_test_lib.MockTempDirTestCase):
-  """Test cros_generate_update_payload.ExtractRoot"""
-  def testTruncate(self):
-    """Test truncating on extraction."""
-    root = os.path.join(self.tempdir, 'root.bin')
-    root_pretruncate = os.path.join(self.tempdir, 'root_pretruncate.bin')
-
-    content = '0123456789'
-    osutils.WriteFile(root, content)
-    self.PatchObject(cros_generate_update_payload, 'ExtractPartitionToTempFile',
-                     return_value=root)
-    self.PatchObject(cros_generate_update_payload, 'Ext2FileSystemSize',
-                     return_value=2)
-
-    cros_generate_update_payload.ExtractRoot(None, root, root_pretruncate)
-    self.assertEqual(osutils.ReadFile(root), content[:2])
-    self.assertEqual(osutils.ReadFile(root_pretruncate), content)
-
-
-class Ext2FileSystemSizeTest(cros_test_lib.MockTestCase):
-  """Test cros_generate_update_payload.Ext2FileSystemSize"""
-  def testExt2FileSystemSize(self):
-    """Test on simple output."""
-    block_size = 4096
-    block_count = 123
-
-    self.PatchObject(cros_build_lib, 'RunCommand',
-                     return_value=cros_build_lib.CommandResult(output='''
-Block size: %d
-Other thing: 123456798
-Not an integer: cdsad132csda
-Block count: %d
-''' % (block_size, block_count)))
-
-    size = cros_generate_update_payload.Ext2FileSystemSize('/dev/null')
-    self.assertEqual(size, block_size * block_count)
-
-
-class ExtractPartitionToTempFileTest(cros_test_lib.MockTempDirTestCase):
-  """Test cros_generate_update_payload.ExtractPartitionToTempFile"""
-  def testExtractPartitionToTempFile(self):
-    """Tests extraction on a simple image."""
-    part_a_bin = '0123'
-    part_b_bin = '4567'
-    image_bin = part_a_bin + part_b_bin
-
-    image = os.path.join(self.tempdir, 'image.bin')
-    osutils.WriteFile(image, image_bin)
-    part_a = os.path.join(self.tempdir, 'a.bin')
-
-    fake_partitions = (
-        cros_build_lib.PartitionInfo(1, 0, 4, 4, 'fs', 'PART-A', ''),
-        cros_build_lib.PartitionInfo(2, 4, 8, 4, 'fs', 'PART-B', ''),
-    )
-    self.PatchObject(cros_build_lib, 'GetImageDiskPartitionInfo',
-                     return_value=fake_partitions)
-
-    cros_generate_update_payload.ExtractPartitionToTempFile(image, 'PART-A',
-                                                            temp_file=part_a)
-    self.assertEqual(osutils.ReadFile(part_a), part_a_bin)
-
-    # Make sure we correctly generate new temp files.
-    tmp = cros_generate_update_payload.ExtractPartitionToTempFile(image,
-                                                                  'PART-B')
-    self.assertEqual(osutils.ReadFile(tmp), part_b_bin)
-
-
 class DeltaGeneratorTest(cros_test_lib.RunCommandTempDirTestCase):
   """Test correct arguments passed to delta_generator."""
   def testDeltaGenerator(self):
     """Test correct arguments propagated to delta_generator call."""
     temp = os.path.join(self.tempdir, 'temp.bin')
     osutils.WriteFile(temp, '0123456789')
-    self.PatchObject(cros_generate_update_payload, 'ExtractPartitionToTempFile',
-                     return_value=temp)
-    self.PatchObject(cros_generate_update_payload, 'Ext2FileSystemSize',
-                     return_value=4)
+    self.PatchObject(partition_lib, 'ExtractPartition', return_value=temp)
+    self.PatchObject(partition_lib, 'Ext2FileSystemSize', return_value=4)
 
     fake_partitions = (
         cros_build_lib.PartitionInfo(3, 0, 4, 4, 'fs', 'ROOT-A', ''),