Repository: google/cloud-forensics-utils

Issue:
Implement a more generic tagging feature for AWS
 [`GetTagForResourceType()`](https://github.com/google/cloud-forensics-utils/blob/master/libcloudforensics/providers/aws/internal/common.py#L35) is currently leveraged to apply only the `Name` tag to resources. This functionality should be more generic and allow a user to provide (e.g. via custom configuration file) which Keys/Values they want to provide in addition to the `Name` tag

Tags tend to be a critical part of organization in an AWS account and help provide clear context and accountability, among other things.

Discussion:
codecov-commenter: # [Codecov](https://codecov.io/gh/google/cloud-forensics-utils/pull/146?src=pr&el=h1) Report
> Merging [#146](https://codecov.io/gh/google/cloud-forensics-utils/pull/146?src=pr&el=desc) into [master](https://codecov.io/gh/google/cloud-forensics-utils/commit/e4448c60002eee03c344ec2803276ff5d91936be&el=desc) will **increase** coverage by `0.49%`.
> The diff coverage is `79.89%`.

[![Impacted file tree graph](https://codecov.io/gh/google/cloud-forensics-utils/pull/146/graphs/tree.svg?width=650&height=150&src=pr&token=zZkPno3ah0)](https://codecov.io/gh/google/cloud-forensics-utils/pull/146?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master     #146      +/-   ##
==========================================
+ Coverage   65.39%   65.88%   +0.49%     
==========================================
  Files          14       14              
  Lines         968     1023      +55     
==========================================
+ Hits          633      674      +41     
- Misses        335      349      +14     
```

| Flag | Coverage Δ | |
|---|---|---|
| #nosetests | `65.88% <79.89%> (+0.49%)` | :arrow_up: |

| [Impacted Files](https://codecov.io/gh/google/cloud-forensics-utils/pull/146?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [libcloudforensics/providers/aws/internal/log.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/146/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2F3cy9pbnRlcm5hbC9sb2cucHk=) | `73.07% <55.55%> (-4.20%)` | :arrow_down: |
| [...ibcloudforensics/providers/gcp/internal/project.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/146/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2djcC9pbnRlcm5hbC9wcm9qZWN0LnB5) | `70.00% <57.14%> (+4.28%)` | :arrow_up: |
| [libcloudforensics/providers/gcp/forensics.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/146/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2djcC9mb3JlbnNpY3MucHk=) | `65.11% <66.66%> (+0.11%)` | :arrow_up: |
| [...cs/providers/gcp/internal/compute\_base\_resource.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/146/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2djcC9pbnRlcm5hbC9jb21wdXRlX2Jhc2VfcmVzb3VyY2UucHk=) | `37.87% <72.22%> (+1.37%)` | :arrow_up: |
| [libcloudforensics/providers/gcp/internal/common.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/146/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2djcC9pbnRlcm5hbC9jb21tb24ucHk=) | `56.04% <72.72%> (-0.28%)` | :arrow_down: |
| [libcloudforensics/providers/aws/forensics.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/146/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2F3cy9mb3JlbnNpY3MucHk=) | `69.56% <80.00%> (-0.21%)` | :arrow_down: |
| [libcloudforensics/providers/aws/internal/ec2.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/146/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2F3cy9pbnRlcm5hbC9lYzIucHk=) | `62.85% <80.00%> (+1.56%)` | :arrow_up: |
| [...bcloudforensics/providers/gcp/internal/function.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/146/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2djcC9pbnRlcm5hbC9mdW5jdGlvbi5weQ==) | `35.29% <83.33%> (+3.03%)` | :arrow_up: |
| [...ibcloudforensics/providers/gcp/internal/compute.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/146/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2djcC9pbnRlcm5hbC9jb21wdXRlLnB5) | `62.97% <85.18%> (+0.62%)` | :arrow_up: |
| [libcloudforensics/providers/gcp/internal/build.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/146/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2djcC9pbnRlcm5hbC9idWlsZC5weQ==) | `33.33% <85.71%> (+3.33%)` | :arrow_up: |
| ... and [4 more](https://codecov.io/gh/google/cloud-forensics-utils/pull/146/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/google/cloud-forensics-utils/pull/146?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/google/cloud-forensics-utils/pull/146?src=pr&el=footer). Last update [ba0561c...25534ed](https://codecov.io/gh/google/cloud-forensics-utils/pull/146?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

giovannt0: Thanks @sa3eed3ed, I think this is a great idea. I've updated the PR, and updated the CLI tool as well. This allows users to call the tool as follow:

```
libcloudforensics aws 'zone' copydisk --volume_id='vol-xxx' --tags='{"customTag": "blah"}'

--

libcloudforensics aws 'zone' copydisk --volume_id='vol-xxx' --tags='{"Name": "some-name", "customTag": "blah"}'
```

In the first case, a volume name will be generated and added as `Name` tag (which is the current behaviour).

In the second case, users can overwrite the name by specifying one on the CLI.

Any other tag added to the dictionary (be it on the CLI or in `CreateVolumeFromSnapshot`, `GetOrCreateAnalysisVm`, `Snapshot` which were all updated) will be added to the AWS resource as well. Note that the `snapshot_name` parameter of the `Snapshot` method was removed as it makes more sense to pass it in the tags directly now.

I've also opened #157 for GCP.

Patch:
diff --git a/examples/aws_cli.py b/examples/aws_cli.py
index d001c4b..d491596 100644
--- a/examples/aws_cli.py
+++ b/examples/aws_cli.py
@@ -17,6 +17,8 @@
 # https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html #pylint: disable=line-too-long
 """Demo CLI tool for AWS."""
 
+import json
+
 from datetime import datetime
 from typing import TYPE_CHECKING
 
@@ -72,7 +74,8 @@ def CreateVolumeCopy(args: 'argparse.Namespace') -> None:
                                            instance_id=args.instance_id,
                                            volume_id=args.volume_id,
                                            src_profile=args.src_profile,
-                                           dst_profile=args.dst_profile)
+                                           dst_profile=args.dst_profile,
+                                           tags=json.loads(args.tags))
   print(
       'Done! Volume {0:s} successfully created. You will find it in '
       'your AWS account under the name {1:s}.'.format(
diff --git a/examples/cli.py b/examples/cli.py
index 309b250..0579602 100644
--- a/examples/cli.py
+++ b/examples/cli.py
@@ -114,7 +114,9 @@ def Main() -> None:
                                   'file.', None),
                 ('--dst_profile', 'The name of the profile for the destination '
                                   'account, as defined in the AWS credentials '
-                                  'file.', None)
+                                  'file.', None),
+                ('--tags', 'A string dictionary of tags to add to the volume '
+                           'copy. ', None)
             ])
   AddParser('aws', aws_subparsers, 'querylogs', 'Query AWS CloudTrail logs',
             args=[
diff --git a/libcloudforensics/providers/aws/forensics.py b/libcloudforensics/providers/aws/forensics.py
index 462e1ec..5bd3573 100644
--- a/libcloudforensics/providers/aws/forensics.py
+++ b/libcloudforensics/providers/aws/forensics.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Forensics on AWS."""
-from typing import TYPE_CHECKING, Tuple, List, Optional
+from typing import TYPE_CHECKING, Tuple, List, Optional, Dict, Any
 
 from libcloudforensics.providers.aws.internal.common import UBUNTU_1804_AMI, LOGGER  # pylint: disable=line-too-long
 from libcloudforensics.providers.aws.internal import account
@@ -27,7 +27,8 @@ def CreateVolumeCopy(zone: str,
                      instance_id: Optional[str] = None,
                      volume_id: Optional[str] = None,
                      src_profile: Optional[str] = None,
-                     dst_profile: Optional[str] = None) -> 'ebs.AWSVolume':
+                     dst_profile: Optional[str] = None,
+                     tags: Optional[Dict[str, str]] = None) -> 'ebs.AWSVolume':
   """Create a copy of an AWS EBS Volume.
 
   By default, the volume copy will be created in the same AWS account where
@@ -83,6 +84,8 @@ def CreateVolumeCopy(zone: str,
     dst_profile (str): Optional. If the volume copy needs to be created in a
         different AWS account, you can specify a different profile name here
         (see example above).
+    tags (Dict[str, str]): Optional. A dictionary of tags to add to the
+          volume copy, for example {'TicketID': 'xxx'}.
 
   Returns:
     AWSVolume: An AWS EBS Volume object.
@@ -136,8 +139,13 @@ def CreateVolumeCopy(zone: str,
       snapshot.aws_account = destination_account
       snapshot = snapshot.Copy(delete=True, deletion_account=source_account)
 
-    new_volume = destination_account.CreateVolumeFromSnapshot(
-        snapshot, volume_name_prefix='evidence')
+    if tags and tags.get('Name'):
+      new_volume = destination_account.CreateVolumeFromSnapshot(
+          snapshot, volume_name=tags['Name'], tags=tags)
+    else:
+      new_volume = destination_account.CreateVolumeFromSnapshot(
+          snapshot, volume_name_prefix='evidence', tags=tags)
+
     snapshot.Delete()
     # Delete the one-time use KMS key, if one was generated
     source_account.DeleteKMSKey(kms_key_id)
@@ -160,7 +168,8 @@ def StartAnalysisVm(
     cpu_cores: int = 4,
     attach_volumes: Optional[List[Tuple[str, str]]] = None,
     dst_profile: Optional[str] = None,
-    ssh_key_name: Optional[str] = None) -> Tuple['ec2.AWSInstance', bool]:
+    ssh_key_name: Optional[str] = None,
+    tags: Optional[Dict[str, str]] = None) -> Tuple['ec2.AWSInstance', bool]:
   """Start a virtual machine for analysis purposes.
 
   Look for an existing AWS instance with tag name vm_name. If found,
@@ -190,6 +199,9 @@ def StartAnalysisVm(
         that if this package fails to install on the target VM, then the VM
         will not be accessible. It is therefore recommended to fill in this
         parameter.
+    tags (Dict[str, str]): Optional. A dictionary of tags to add to the
+          instance, for example {'TicketID': 'xxx'}. An entry for the instance
+          name is added by default.
 
   Returns:
     Tuple[AWSInstance, bool]: a tuple with a virtual machine object
@@ -202,7 +214,8 @@ def StartAnalysisVm(
       boot_volume_size,
       ami,
       cpu_cores,
-      ssh_key_name=ssh_key_name)
+      ssh_key_name=ssh_key_name,
+      tags=tags)
   for volume_id, device_name in (attach_volumes or []):
     analysis_vm.AttachVolume(aws_account.GetVolumeById(volume_id), device_name)
   return analysis_vm, created
diff --git a/libcloudforensics/providers/aws/internal/account.py b/libcloudforensics/providers/aws/internal/account.py
index 1d8db9f..9d1b598 100644
--- a/libcloudforensics/providers/aws/internal/account.py
+++ b/libcloudforensics/providers/aws/internal/account.py
@@ -399,7 +399,8 @@ class AWSAccount:
       snapshot: ebs.AWSSnapshot,
       volume_name: Optional[str] = None,
       volume_name_prefix: str = '',
-      kms_key_id: Optional[str] = None) -> ebs.AWSVolume:
+      kms_key_id: Optional[str] = None,
+      tags: Optional[Dict[str, str]] = None) -> ebs.AWSVolume:
     """Create a new volume based on a snapshot.
 
     Args:
@@ -407,6 +408,9 @@ class AWSAccount:
       volume_name (str): Optional. String to use as new volume name.
       volume_name_prefix (str): Optional. String to prefix the volume name with.
       kms_key_id (str): Optional. A KMS key id to encrypt the volume with.
+      tags (Dict[str, str]): Optional. A dictionary of tags to add to the
+          volume, for example {'TicketID': 'xxx'}. An entry for the volume
+          name is added by default.
 
     Returns:
       AWSVolume: An AWS EBS Volume.
@@ -425,12 +429,15 @@ class AWSAccount:
           'Volume name {0:s} does not comply with '
           '{1:s}'.format(volume_name, common.REGEX_TAG_VALUE.pattern))
 
+    if not tags:
+      tags = {}
+    tags['Name'] = volume_name
+
     client = self.ClientApi(common.EC2_SERVICE)
     create_volume_args = {
         'AvailabilityZone': snapshot.availability_zone,
         'SnapshotId': snapshot.snapshot_id,
-        'TagSpecifications':
-            [common.GetTagForResourceType('volume', volume_name)]
+        'TagSpecifications': [common.CreateTags(common.VOLUME, tags)]
     }
     if kms_key_id:
       create_volume_args['Encrypted'] = True
@@ -463,7 +470,8 @@ class AWSAccount:
       ami: str,
       cpu_cores: int,
       packages: Optional[List[str]] = None,
-      ssh_key_name: Optional[str] = None) -> Tuple[ec2.AWSInstance, bool]:
+      ssh_key_name: Optional[str] = None,
+      tags: Optional[Dict[str, str]] = None) -> Tuple[ec2.AWSInstance, bool]:
     """Get or create a new virtual machine for analysis purposes.
 
     Args:
@@ -479,6 +487,9 @@ class AWSAccount:
           that if this package fails to install on the target VM, then the VM
           will not be accessible. It is therefore recommended to fill in this
           parameter.
+      tags (Dict[str, str]): Optional. A dictionary of tags to add to the
+          instance, for example {'TicketID': 'xxx'}. An entry for the instance
+          name is added by default.
 
     Returns:
       Tuple[AWSInstance, bool]: A tuple with an AWSInstance object and a
@@ -509,6 +520,10 @@ class AWSAccount:
         '(exit ${exit_code})',
         'apt -y install ec2-instance-connect && (exit ${exit_code})')
 
+    if not tags:
+      tags = {}
+    tags['Name'] = vm_name
+
     client = self.ClientApi(common.EC2_SERVICE)
     vm_args = {
         'BlockDeviceMappings':
@@ -517,8 +532,7 @@ class AWSAccount:
         'MinCount': 1,
         'MaxCount': 1,
         'InstanceType': instance_type,
-        'TagSpecifications':
-            [common.GetTagForResourceType('instance', vm_name)],
+        'TagSpecifications': [common.CreateTags(common.INSTANCE, tags)],
         'UserData': startup_script,
         'Placement': {'AvailabilityZone': self.default_availability_zone}
     }
diff --git a/libcloudforensics/providers/aws/internal/common.py b/libcloudforensics/providers/aws/internal/common.py
index c911805..cac43aa 100644
--- a/libcloudforensics/providers/aws/internal/common.py
+++ b/libcloudforensics/providers/aws/internal/common.py
@@ -25,6 +25,11 @@ ACCOUNT_SERVICE = 'sts'
 KMS_SERVICE = 'kms'
 CLOUDTRAIL_SERVICE = 'cloudtrail'
 
+# Resource types constant
+INSTANCE = 'instance'
+VOLUME = 'volume'
+SNAPSHOT = 'snapshot'
+
 # Default Amazon Machine Image to use for bootstrapping instances
 UBUNTU_1804_AMI = 'ami-0013b3aa57f8a4331'
 REGEX_TAG_VALUE = re.compile('^.{1,255}$')
@@ -32,26 +37,27 @@ REGEX_TAG_VALUE = re.compile('^.{1,255}$')
 LOGGER = logging.getLogger()
 
 
-def GetTagForResourceType(resource: str, name: str) -> Dict[str, Any]:
-  """Create a dictionary for AWS Tag Specifications.
+def CreateTags(resource: str, tags: Dict[str, str]) -> Dict[str, Any]:
+  """Create AWS Tag Specifications.
 
   Args:
     resource (str): The type of AWS resource.
-    name (str): The name of the resource.
+    tags (Dict[str, str]): A dictionary of tags to add to the resource.
 
   Returns:
-    Dict[str, str|List[Dict]]: A dictionary for AWS Tag Specifications.
+    Dict[str, Any]: A dictionary for AWS Tag Specifications.
   """
 
-  return {
+  tag_specifications = {
       'ResourceType': resource,
-      'Tags': [
-          {
-              'Key': 'Name',
-              'Value': name
-          }
-      ]
-  }
+      'Tags': []
+  }  # type: Dict[str, Any]
+  for tag in tags:
+    tag_specifications['Tags'].append({
+        'Key': tag,
+        'Value': tags[tag]
+    })
+  return tag_specifications
 
 
 def GetInstanceTypeByCPU(cpu_cores: int) -> str:
diff --git a/libcloudforensics/providers/aws/internal/ebs.py b/libcloudforensics/providers/aws/internal/ebs.py
index 1d60ba6..7937ab6 100644
--- a/libcloudforensics/providers/aws/internal/ebs.py
+++ b/libcloudforensics/providers/aws/internal/ebs.py
@@ -14,7 +14,6 @@
 # limitations under the License.
 """Disk functionality."""
 
-from datetime import datetime
 from typing import TYPE_CHECKING, Dict, Optional, Union
 
 import botocore
@@ -105,11 +104,12 @@ class AWSVolume(AWSElasticBlockStore):
     self.volume_id = volume_id
     self.device_name = device_name
 
-  def Snapshot(self, snapshot_name: Optional[str] = None) -> 'AWSSnapshot':
+  def Snapshot(self, tags: Optional[Dict[str, str]] = None) -> 'AWSSnapshot':
     """Create a snapshot of the volume.
 
     Args:
-      snapshot_name (str): Optional. Name tag of the snapshot.
+      tags (Dict[str, str]): Optional. A dictionary of tags to add to the
+          snapshot, for example {'Name': 'my-snapshot-name', 'TicketID': 'xxx'}.
 
     Returns:
       AWSSnapshot: A snapshot object.
@@ -119,22 +119,23 @@ class AWSVolume(AWSElasticBlockStore):
       RuntimeError: If the snapshot could not be created.
     """
 
-    timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
-    if not snapshot_name:
-      snapshot_name = self.volume_id
-    truncate_at = 255 - len(timestamp) - 1
-    snapshot_name = '{0}-{1}'.format(snapshot_name[:truncate_at], timestamp)
+    if not tags:
+      tags = {}
+
+    snapshot_name = tags.get('Name') or (self.volume_id + '-snapshot')
+    truncate_at = 255 - 1
+    snapshot_name = snapshot_name[:truncate_at]
     if not common.REGEX_TAG_VALUE.match(snapshot_name):
       raise ValueError('Snapshot name {0:s} does not comply with '
                        '{1:s}'.format(snapshot_name,
                                       common.REGEX_TAG_VALUE.pattern))
+    tags['Name'] = snapshot_name
 
     client = self.aws_account.ClientApi(common.EC2_SERVICE)
     try:
       snapshot = client.create_snapshot(
           VolumeId=self.volume_id,
-          TagSpecifications=[common.GetTagForResourceType(
-              'snapshot', snapshot_name)])
+          TagSpecifications=[common.CreateTags(common.SNAPSHOT, tags)])
 
       snapshot_id = snapshot.get('SnapshotId')
       # Wait for snapshot completion
diff --git a/libcloudforensics/providers/aws/internal/ec2.py b/libcloudforensics/providers/aws/internal/ec2.py
index f8d9bc9..e4f1a69 100644
--- a/libcloudforensics/providers/aws/internal/ec2.py
+++ b/libcloudforensics/providers/aws/internal/ec2.py
@@ -15,7 +15,7 @@
 """Instance functionality."""
 from typing import TYPE_CHECKING, Dict, Optional
 
-from libcloudforensics.providers.aws.internal.common import EC2_SERVICE
+from libcloudforensics.providers.aws.internal import common
 
 if TYPE_CHECKING:
   # TYPE_CHECKING is always False at runtime, therefore it is safe to ignore
@@ -69,7 +69,7 @@ class AWSInstance:
     """
 
     boot_device = self.aws_account.ResourceApi(
-        EC2_SERVICE).Instance(self.instance_id).root_device_name
+        common.EC2_SERVICE).Instance(self.instance_id).root_device_name
     volumes = self.ListVolumes()
 
     for volume_id in volumes:
@@ -127,7 +127,7 @@ class AWSInstance:
       RuntimeError: If the volume could not be attached.
     """
 
-    client = self.aws_account.ClientApi(EC2_SERVICE)
+    client = self.aws_account.ClientApi(common.EC2_SERVICE)
     try:
       client.attach_volume(Device=device_name,
                            InstanceId=self.instance_id,


Test patch:
diff --git a/tests/providers/aws/aws_test.py b/tests/providers/aws/aws_test.py
index 0582097..ee9656f 100644
--- a/tests/providers/aws/aws_test.py
+++ b/tests/providers/aws/aws_test.py
@@ -425,16 +425,6 @@ class AWSAccountTest(unittest.TestCase):
     # pylint: enable=protected-access
     self.assertEqual(50, config['Ebs']['VolumeSize'])
 
-  @typing.no_type_check
-  def testGetInstanceTypeByCPU(self):
-    """Test that the instance type matches the requested amount of CPU cores."""
-    # pylint: disable=protected-access
-    self.assertEqual('m4.large', common.GetInstanceTypeByCPU(2))
-    self.assertEqual('m4.16xlarge', common.GetInstanceTypeByCPU(64))
-    self.assertRaises(ValueError, common.GetInstanceTypeByCPU, 0)
-    self.assertRaises(ValueError, common.GetInstanceTypeByCPU, 256)
-    # pylint: enable=protected-access
-
 
 class AWSInstanceTest(unittest.TestCase):
   """Test AWSInstance class."""
@@ -467,19 +457,15 @@ class AWSVolumeTest(unittest.TestCase):
     snapshot.return_value = MOCK_CREATE_SNAPSHOT
     mock_ec2_api.return_value.get_waiter.return_value.wait.return_value = None
 
-    # Snapshot(snapshot_name=None). Snapshot should start with the volume's name
+    # Snapshot(tags=None). Snapshot name should be the volume's id.
     snapshot = FAKE_VOLUME.Snapshot()
     self.assertIsInstance(snapshot, ebs.AWSSnapshot)
-    # Part of the snapshot name is taken from a timestamp, therefore we only
-    # assert for the beginning of the string.
-    self.assertTrue(snapshot.name.startswith('fake-volume'))
+    self.assertEqual('fake-volume-id-snapshot', snapshot.name)
 
-    # Snapshot(snapshot_name='my-Snapshot'). Snapshot should start with
-    # 'my-Snapshot'
-    snapshot = FAKE_VOLUME.Snapshot(snapshot_name='my-snapshot')
+    # Snapshot(tags={'Name': 'my-snapshot'}). Snapshot should be 'my-snapshot'.
+    snapshot = FAKE_VOLUME.Snapshot(tags={'Name': 'my-snapshot'})
     self.assertIsInstance(snapshot, ebs.AWSSnapshot)
-    # Same as above regarding the timestamp.
-    self.assertTrue(snapshot.name.startswith('my-snapshot'))
+    self.assertEqual('my-snapshot', snapshot.name)
 
 
 class AWSSnapshotTest(unittest.TestCase):
@@ -519,8 +505,35 @@ class AWSCloudTrailTest(unittest.TestCase):
     self.assertEqual(FAKE_EVENT_LIST[0], lookup_events[0])
 
 
+class AWSCommon(unittest.TestCase):
+  """Test the common.py public methods"""
+
+  @typing.no_type_check
+  def testCreateTags(self):
+    """Test that tag specifications are correclty created"""
+    tag_specifications = common.CreateTags(common.VOLUME, {'Name': 'fake-name'})
+    self.assertEqual('volume', tag_specifications['ResourceType'])
+    self.assertEqual(1, len(tag_specifications['Tags']))
+    self.assertEqual('Name', tag_specifications['Tags'][0]['Key'])
+    self.assertEqual('fake-name', tag_specifications['Tags'][0]['Value'])
+
+    tag_specifications = common.CreateTags(
+        common.VOLUME, {'Name': 'fake-name', 'FakeTag': 'fake-tag'})
+    self.assertEqual(2, len(tag_specifications['Tags']))
+    self.assertEqual('FakeTag', tag_specifications['Tags'][1]['Key'])
+    self.assertEqual('fake-tag', tag_specifications['Tags'][1]['Value'])
+
+  @typing.no_type_check
+  def testGetInstanceTypeByCPU(self):
+    """Test that the instance type matches the requested amount of CPU cores."""
+    self.assertEqual('m4.large', common.GetInstanceTypeByCPU(2))
+    self.assertEqual('m4.16xlarge', common.GetInstanceTypeByCPU(64))
+    self.assertRaises(ValueError, common.GetInstanceTypeByCPU, 0)
+    self.assertRaises(ValueError, common.GetInstanceTypeByCPU, 256)
+
+
 class AWSTest(unittest.TestCase):
-  """Test the account.py public methods."""
+  """Test the forensics.py public methods."""
   # pylint: disable=line-too-long
 
   @typing.no_type_check